hayson-engine 0.22.0

use std::collections::HashMap;

use serde_json::Value;

use crate::registry::table_registry::TableRegistry;
use crate::schema::HaysonRecord;

use super::planner::{execute_query, QuerySpec, RecordSource, RecordSourceError};

#[derive(Debug, Clone, Default, PartialEq)]
pub struct AggregateSpec {
    pub count: bool,
    pub sum: Vec<String>,
    pub avg: Vec<String>,
    pub min: Vec<String>,
    pub max: Vec<String>,
    pub group_by: Option<String>,
}

#[derive(Debug, Clone, Default, PartialEq)]
pub struct AggregateResult {
    pub count: usize,
    pub sum: HashMap<String, f64>,
    pub avg: HashMap<String, Option<f64>>,
    pub min: HashMap<String, Option<f64>>,
    pub max: HashMap<String, Option<f64>>,
}

/// One `AggregateResult` per distinct value of the `group_by` field, sorted by that value so
/// output order is deterministic. Records missing the field, or holding a non-scalar value
/// (object/array), fall into a single `Value::Null` group.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct GroupedAggregateResult(pub Vec<(Value, AggregateResult)>);

struct Accumulator {
    sum: f64,
    seen: usize,
    min: Option<f64>,
    max: Option<f64>,
}

impl Accumulator {
    fn new() -> Self {
        Accumulator {
            sum: 0.0,
            seen: 0,
            min: None,
            max: None,
        }
    }

    fn add(&mut self, value: f64) {
        self.sum += value;
        self.seen += 1;
        self.min = Some(self.min.map_or(value, |current| current.min(value)));
        self.max = Some(self.max.map_or(value, |current| current.max(value)));
    }
}

fn numeric_field(record: &HaysonRecord, field: &str) -> Option<f64> {
    match record.get(field) {
        Some(Value::Number(number)) => number.as_f64(),
        _ => None,
    }
}

/// Runs `spec` through `execute_query` — unpaged, because `limit`/`offset` window the answer, not
/// the rows an aggregate folds — then for every field named by `sum`/`avg`/`min`/`max`
/// re-reads the matched records through `source` and accumulates, skipping a record for a given
/// field whenever that field is missing or not a JSON number. `count` always reflects every
/// matched key regardless of field types. An empty match (or a field with no numeric values at
/// all) reports `sum: 0.0`, `avg`/`min`/`max: None` — a proven-empty aggregate, not a missing one.
pub fn execute_aggregate<S: RecordSource + ?Sized>(
    registry: &TableRegistry,
    source: &mut S,
    spec: &QuerySpec,
    aggregate: &AggregateSpec,
) -> Result<AggregateResult, RecordSourceError> {
    let matched = execute_query(registry, source, &spec.unpaged())?;

    let mut accumulators: HashMap<&str, Accumulator> = HashMap::new();
    for field in aggregate
        .sum
        .iter()
        .chain(&aggregate.avg)
        .chain(&aggregate.min)
        .chain(&aggregate.max)
    {
        accumulators.entry(field.as_str()).or_insert_with(Accumulator::new);
    }

    if !accumulators.is_empty() {
        for key in &matched.keys {
            let Some(record) = source.load(key)? else { continue };
            for (field, accumulator) in accumulators.iter_mut() {
                if let Some(value) = numeric_field(&record, field) {
                    accumulator.add(value);
                }
            }
        }
    }

    let mut result = AggregateResult {
        count: matched.keys.len(),
        ..Default::default()
    };
    for field in &aggregate.sum {
        result.sum.insert(field.clone(), accumulators[field.as_str()].sum);
    }
    for field in &aggregate.avg {
        let accumulator = &accumulators[field.as_str()];
        result.avg.insert(
            field.clone(),
            (accumulator.seen > 0).then(|| accumulator.sum / accumulator.seen as f64),
        );
    }
    for field in &aggregate.min {
        result.min.insert(field.clone(), accumulators[field.as_str()].min);
    }
    for field in &aggregate.max {
        result.max.insert(field.clone(), accumulators[field.as_str()].max);
    }

    Ok(result)
}

fn group_key_value(record: &HaysonRecord, field: &str) -> Value {
    match record.get(field) {
        Some(Value::Object(_)) | Some(Value::Array(_)) | None => Value::Null,
        Some(value) => value.clone(),
    }
}

fn compare_group_keys(a: &Value, b: &Value) -> std::cmp::Ordering {
    use std::cmp::Ordering;

    fn rank(value: &Value) -> u8 {
        match value {
            Value::Null => 0,
            Value::Bool(_) => 1,
            Value::Number(_) => 2,
            Value::String(_) => 3,
            _ => 4,
        }
    }

    match rank(a).cmp(&rank(b)) {
        Ordering::Equal => match (a, b) {
            (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
            (Value::Number(x), Value::Number(y)) => x.as_f64().partial_cmp(&y.as_f64()).unwrap_or(Ordering::Equal),
            (Value::String(x), Value::String(y)) => x.cmp(y),
            _ => a.to_string().cmp(&b.to_string()),
        },
        other => other,
    }
}

struct GroupState {
    count: usize,
    accumulators: HashMap<String, Accumulator>,
}

/// Same accumulation as `execute_aggregate`, run independently per distinct value of
/// `aggregate.group_by` (required — panics if unset). Groups are emitted sorted by that value.
pub fn execute_grouped_aggregate<S: RecordSource + ?Sized>(
    registry: &TableRegistry,
    source: &mut S,
    spec: &QuerySpec,
    aggregate: &AggregateSpec,
) -> Result<GroupedAggregateResult, RecordSourceError> {
    let group_field = aggregate
        .group_by
        .as_deref()
        .expect("execute_grouped_aggregate requires group_by");
    let matched = execute_query(registry, source, &spec.unpaged())?;

    let numeric_fields: Vec<&str> = aggregate
        .sum
        .iter()
        .chain(&aggregate.avg)
        .chain(&aggregate.min)
        .chain(&aggregate.max)
        .map(String::as_str)
        .collect();

    let mut groups: HashMap<Value, GroupState> = HashMap::new();
    for key in &matched.keys {
        let Some(record) = source.load(key)? else { continue };
        let group_key = group_key_value(&record, group_field);
        let state = groups.entry(group_key).or_insert_with(|| GroupState {
            count: 0,
            accumulators: numeric_fields
                .iter()
                .map(|field| (field.to_string(), Accumulator::new()))
                .collect(),
        });
        state.count += 1;
        for field in &numeric_fields {
            if let Some(value) = numeric_field(&record, field) {
                state.accumulators.get_mut(*field).unwrap().add(value);
            }
        }
    }

    let mut entries: Vec<(Value, AggregateResult)> = groups
        .into_iter()
        .map(|(key, state)| {
            let mut result = AggregateResult {
                count: state.count,
                ..Default::default()
            };
            for field in &aggregate.sum {
                result.sum.insert(field.clone(), state.accumulators[field.as_str()].sum);
            }
            for field in &aggregate.avg {
                let accumulator = &state.accumulators[field.as_str()];
                result.avg.insert(
                    field.clone(),
                    (accumulator.seen > 0).then(|| accumulator.sum / accumulator.seen as f64),
                );
            }
            for field in &aggregate.min {
                result.min.insert(field.clone(), state.accumulators[field.as_str()].min);
            }
            for field in &aggregate.max {
                result.max.insert(field.clone(), state.accumulators[field.as_str()].max);
            }
            (key, result)
        })
        .collect();

    entries.sort_by(|(a, _), (b, _)| compare_group_keys(a, b));
    Ok(GroupedAggregateResult(entries))
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;
    use crate::query::builder::QueryBuilder;
    use crate::query::planner::LogRecordSource;
    use crate::query::predicates::ComparisonOp;
    use crate::schema::parse_schema;
    use crate::storage::record::{serialize_entry, LogMeta, LogOp};
    use crate::storage::LogFile;

    struct Fixture {
        _dir: tempfile::TempDir,
        registry: TableRegistry,
        log: LogFile,
    }

    fn schema() -> crate::schema::TableSchema {
        parse_schema(&json!({
            "schemaVersion": 1,
            "table": "orders",
            "primaryKey": "id",
            "fields": {
                "id": { "type": "string" },
                "status": { "type": "string", "index": "exact" },
                "amount": { "type": "float" },
            },
        }))
        .unwrap()
    }

    fn record(pairs: &[(&str, Value)]) -> HaysonRecord {
        pairs.iter().map(|(k, v)| (k.to_string(), v.clone())).collect()
    }

    fn fixture(records: &[HaysonRecord]) -> Fixture {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("orders.jsonl");
        let mut log = LogFile::open(&path, false).unwrap();
        let mut registry = TableRegistry::new(schema());

        let mut seq = 0u64;
        for entry in records {
            seq += 1;
            let key = entry.get("id").and_then(Value::as_str).unwrap().to_string();
            let line = serialize_entry(
                &LogMeta {
                    key: key.clone(),
                    op: LogOp::Put,
                    seq,
                    ts: seq as i64,
                },
                Some(entry),
            );
            let slots = log.append(&[line]).unwrap();
            registry.apply_put(&key, entry, slots[0], seq);
        }
        Fixture {
            _dir: dir,
            registry,
            log,
        }
    }

    fn run(fixture: &Fixture, spec: &QuerySpec, aggregate: &AggregateSpec) -> AggregateResult {
        let mut source = LogRecordSource::new(&fixture.registry, &fixture.log);
        execute_aggregate(&fixture.registry, &mut source, spec, aggregate).unwrap()
    }

    #[test]
    fn counts_every_matching_record() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(20.0))]),
            record(&[
                ("id", json!("o3")),
                ("status", json!("pending")),
                ("amount", json!(5.0)),
            ]),
        ]);
        let spec = QueryBuilder::new()
            .r#where("status", ComparisonOp::Eq, json!("paid"))
            .spec();
        let aggregate = AggregateSpec {
            count: true,
            ..Default::default()
        };

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.count, 2);
    }

    #[test]
    fn sum_avg_min_max_over_a_numeric_field() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(30.0))]),
            record(&[("id", json!("o3")), ("status", json!("paid")), ("amount", json!(20.0))]),
        ]);
        let spec = QueryBuilder::new()
            .r#where("status", ComparisonOp::Eq, json!("paid"))
            .spec();
        let aggregate = AggregateSpec {
            sum: vec!["amount".into()],
            avg: vec!["amount".into()],
            min: vec!["amount".into()],
            max: vec!["amount".into()],
            ..Default::default()
        };

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.sum["amount"], 60.0);
        assert_eq!(result.avg["amount"], Some(20.0));
        assert_eq!(result.min["amount"], Some(10.0));
        assert_eq!(result.max["amount"], Some(30.0));
    }

    #[test]
    fn non_numeric_values_are_skipped_without_panicking() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!("n/a"))]),
            record(&[("id", json!("o3")), ("status", json!("paid")), ("amount", json!(30.0))]),
        ]);
        let spec = QueryBuilder::new()
            .r#where("status", ComparisonOp::Eq, json!("paid"))
            .spec();
        let aggregate = AggregateSpec {
            count: true,
            sum: vec!["amount".into()],
            avg: vec!["amount".into()],
            ..Default::default()
        };

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.count, 3, "count ignores field types entirely");
        assert_eq!(result.sum["amount"], 40.0);
        assert_eq!(result.avg["amount"], Some(20.0));
    }

    #[test]
    fn an_empty_match_reports_zero_count_and_none_for_the_rest() {
        let fixture = fixture(&[record(&[
            ("id", json!("o1")),
            ("status", json!("paid")),
            ("amount", json!(10.0)),
        ])]);
        let spec = QueryBuilder::new()
            .r#where("status", ComparisonOp::Eq, json!("ghost"))
            .spec();
        let aggregate = AggregateSpec {
            count: true,
            sum: vec!["amount".into()],
            avg: vec!["amount".into()],
            min: vec!["amount".into()],
            max: vec!["amount".into()],
            ..Default::default()
        };

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.count, 0);
        assert_eq!(result.sum["amount"], 0.0);
        assert_eq!(result.avg["amount"], None);
        assert_eq!(result.min["amount"], None);
        assert_eq!(result.max["amount"], None);
    }

    #[test]
    fn a_field_with_no_numeric_values_at_all_behaves_like_an_empty_aggregate() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!("n/a"))]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!("n/a"))]),
        ]);
        let spec = QueryBuilder::new()
            .r#where("status", ComparisonOp::Eq, json!("paid"))
            .spec();
        let aggregate = AggregateSpec {
            sum: vec!["amount".into()],
            avg: vec!["amount".into()],
            ..Default::default()
        };

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.sum["amount"], 0.0);
        assert_eq!(result.avg["amount"], None);
    }

    #[test]
    fn builder_wires_the_aggregate_methods_into_a_spec_pair() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(30.0))]),
        ]);
        let (spec, aggregate) = QueryBuilder::new()
            .r#where("status", ComparisonOp::Eq, json!("paid"))
            .count()
            .sum("amount")
            .avg("amount")
            .build_aggregate();

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.count, 2);
        assert_eq!(result.sum["amount"], 40.0);
        assert_eq!(result.avg["amount"], Some(20.0));
    }

    fn run_grouped(fixture: &Fixture, spec: &QuerySpec, aggregate: &AggregateSpec) -> GroupedAggregateResult {
        let mut source = LogRecordSource::new(&fixture.registry, &fixture.log);
        execute_grouped_aggregate(&fixture.registry, &mut source, spec, aggregate).unwrap()
    }

    #[test]
    fn groups_count_and_sum_independently_per_distinct_field_value() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(30.0))]),
            record(&[
                ("id", json!("o3")),
                ("status", json!("pending")),
                ("amount", json!(5.0)),
            ]),
        ]);
        let (spec, aggregate) = QueryBuilder::new()
            .count()
            .sum("amount")
            .group_by("status")
            .build_aggregate();

        let GroupedAggregateResult(groups) = run_grouped(&fixture, &spec, &aggregate);
        assert_eq!(groups.len(), 2);

        let (paid_key, paid) = &groups[0];
        assert_eq!(paid_key, &json!("paid"));
        assert_eq!(paid.count, 2);
        assert_eq!(paid.sum["amount"], 40.0);

        let (pending_key, pending) = &groups[1];
        assert_eq!(pending_key, &json!("pending"));
        assert_eq!(pending.count, 1);
        assert_eq!(pending.sum["amount"], 5.0);
    }

    #[test]
    fn records_missing_the_group_field_fall_into_a_null_group() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[("id", json!("o2")), ("amount", json!(15.0))]),
        ]);
        let (spec, aggregate) = QueryBuilder::new().count().group_by("status").build_aggregate();

        let GroupedAggregateResult(groups) = run_grouped(&fixture, &spec, &aggregate);
        assert_eq!(groups.len(), 2);

        let null_group = groups.iter().find(|(key, _)| key.is_null()).unwrap();
        assert_eq!(null_group.1.count, 1);
    }

    #[test]
    fn grouped_output_is_sorted_deterministically_by_group_key() {
        let fixture = fixture(&[
            record(&[
                ("id", json!("o1")),
                ("status", json!("shipped")),
                ("amount", json!(1.0)),
            ]),
            record(&[("id", json!("o2")), ("status", json!("paid")), ("amount", json!(1.0))]),
            record(&[("id", json!("o3")), ("amount", json!(1.0))]),
            record(&[
                ("id", json!("o4")),
                ("status", json!("pending")),
                ("amount", json!(1.0)),
            ]),
        ]);
        let (spec, aggregate) = QueryBuilder::new().count().group_by("status").build_aggregate();

        let GroupedAggregateResult(groups) = run_grouped(&fixture, &spec, &aggregate);
        let keys: Vec<Value> = groups.into_iter().map(|(key, _)| key).collect();
        assert_eq!(
            keys,
            vec![json!(null), json!("paid"), json!("pending"), json!("shipped")]
        );
    }

    #[test]
    fn ungrouped_aggregation_is_unaffected_by_the_group_by_field_existing() {
        let fixture = fixture(&[
            record(&[("id", json!("o1")), ("status", json!("paid")), ("amount", json!(10.0))]),
            record(&[
                ("id", json!("o2")),
                ("status", json!("pending")),
                ("amount", json!(5.0)),
            ]),
        ]);
        let spec = QueryBuilder::new().spec();
        let aggregate = AggregateSpec {
            count: true,
            sum: vec!["amount".into()],
            ..Default::default()
        };

        let result = run(&fixture, &spec, &aggregate);
        assert_eq!(result.count, 2);
        assert_eq!(result.sum["amount"], 15.0);
    }
}