reddb-io-server 1.24.1

//! Per-group accumulator state.
//!
//! Each group owns one [`GroupAccumulator`]. `accumulate` folds a
//! single [`super::ScanRow`] into the slot vector; `finalize` emits
//! the per-position [`Value`]s the planner returns.
//!
//! No per-row materialisation lives here — the slot layout is sized
//! once at construction time from the AST and stays constant for
//! the life of the group.

use super::ast::{AggregateExpr, AggregateOp};
use reddb_types::Value;

#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};

/// Test-only counter incremented whenever the planner emits a
/// final per-group row. The legacy path materialises one row per
/// scanned input; this path materialises one per *group*. The
/// gap is the whole point of the optimisation.
#[cfg(test)]
pub(crate) static MATERIALIZED_COUNT: AtomicUsize = AtomicUsize::new(0);

#[cfg(test)]
pub(crate) fn note_materialized() {
    MATERIALIZED_COUNT.fetch_add(1, Ordering::Relaxed);
}

#[cfg(not(test))]
#[inline]
pub(crate) fn note_materialized() {}

/// Per-aggregate accumulator slot.
///
/// One slot per [`AggregateExpr`] in the plan. Variants are picked
/// by [`AggregateOp`] at construction; the hot loop matches on the
/// variant rather than the op so each branch has a stable layout.
enum Slot {
    /// `COUNT(*)` — every accumulated row bumps the counter.
    CountStar { count: u64 },
    /// `COUNT(col)` — non-NULL inputs only.
    CountColumn { count: u64 },
    /// `SUM(col)` — running f64 sum, ignoring NULLs. We track
    /// `seen_any` so we can return NULL for an entirely-NULL group
    /// (matching SQL semantics) instead of `0.0`.
    Sum { sum: f64, seen_any: bool },
    /// `AVG(col)` — `(sum, count)` finalised at emission. Same
    /// "all-NULL → NULL" semantics as SUM.
    Avg { sum: f64, count: u64 },
    /// `MIN(col)` — running extremum, type-preserving. We compare
    /// via the canonical key encoding so types stay consistent
    /// across the scan (no implicit numeric coercion that the
    /// legacy path also avoids).
    Min { current: Option<Value> },
    /// `MAX(col)` — symmetric.
    Max { current: Option<Value> },
}

impl Slot {
    fn for_op(op: AggregateOp) -> Self {
        match op {
            AggregateOp::CountStar => Slot::CountStar { count: 0 },
            AggregateOp::CountColumn => Slot::CountColumn { count: 0 },
            AggregateOp::Sum => Slot::Sum {
                sum: 0.0,
                seen_any: false,
            },
            AggregateOp::Avg => Slot::Avg { sum: 0.0, count: 0 },
            AggregateOp::Min => Slot::Min { current: None },
            AggregateOp::Max => Slot::Max { current: None },
        }
    }
}

/// One per group. Owns one [`Slot`] per aggregate expression, in
/// AST order.
pub(super) struct GroupAccumulator {
    slots: Vec<Slot>,
}

impl GroupAccumulator {
    pub(super) fn new(aggregates: &[AggregateExpr]) -> Self {
        Self {
            slots: aggregates.iter().map(|a| Slot::for_op(a.op)).collect(),
        }
    }

    /// Fold one scan row into every slot.
    ///
    /// Per-aggregate input is read by index; out-of-range indices
    /// are surfaced through the plan-time check in
    /// [`super::AggregateQueryPlanner::plan`] and never reach the
    /// hot loop.
    pub(super) fn accumulate(&mut self, aggregates: &[AggregateExpr], inputs: &[Value]) {
        for (slot, expr) in self.slots.iter_mut().zip(aggregates.iter()) {
            match slot {
                Slot::CountStar { count } => {
                    *count += 1;
                }
                Slot::CountColumn { count } => {
                    if let Some(v) = inputs.get(expr.input_index) {
                        if !matches!(v, Value::Null) {
                            *count += 1;
                        }
                    }
                }
                Slot::Sum { sum, seen_any } => {
                    if let Some(v) = inputs.get(expr.input_index) {
                        if let Some(n) = super::super::aggregate_value_to_f64(v) {
                            *sum += n;
                            *seen_any = true;
                        }
                    }
                }
                Slot::Avg { sum, count } => {
                    if let Some(v) = inputs.get(expr.input_index) {
                        if let Some(n) = super::super::aggregate_value_to_f64(v) {
                            *sum += n;
                            *count += 1;
                        }
                    }
                }
                Slot::Min { current } => {
                    if let Some(v) = inputs.get(expr.input_index) {
                        update_extreme(current, v, std::cmp::Ordering::Less);
                    }
                }
                Slot::Max { current } => {
                    if let Some(v) = inputs.get(expr.input_index) {
                        update_extreme(current, v, std::cmp::Ordering::Greater);
                    }
                }
            }
        }
    }

    /// Emit the per-aggregate result row. One [`Value`] per slot,
    /// in AST order. Bumps [`MATERIALIZED_COUNT`] in tests so the
    /// "O(group count)" invariant can be asserted.
    pub(super) fn finalize(self) -> Vec<Value> {
        note_materialized();
        self.slots
            .into_iter()
            .map(|slot| match slot {
                Slot::CountStar { count } | Slot::CountColumn { count } => {
                    Value::Integer(count as i64)
                }
                Slot::Sum { sum, seen_any } => {
                    if seen_any {
                        sum_f64_to_value(sum)
                    } else {
                        Value::Null
                    }
                }
                Slot::Avg { sum, count } => {
                    if count == 0 {
                        Value::Null
                    } else {
                        Value::Float(sum / count as f64)
                    }
                }
                Slot::Min { current } | Slot::Max { current } => current.unwrap_or(Value::Null),
            })
            .collect()
    }
}

fn sum_f64_to_value(f: f64) -> Value {
    if f.fract() == 0.0 && f >= i64::MIN as f64 && f <= i64::MAX as f64 {
        Value::Integer(f as i64)
    } else {
        Value::Float(f)
    }
}

/// Update `current` if `candidate` extends `target_ordering`.
///
/// `Ordering::Less` → MIN behaviour (replace when candidate < current).
/// `Ordering::Greater` → MAX behaviour.
///
/// NULL inputs are skipped. Non-comparable pairs (different kinds,
/// non-finite floats) leave `current` untouched — same conservative
/// rule the legacy path uses, since SQL doesn't define an order
/// across families.
fn update_extreme(current: &mut Option<Value>, candidate: &Value, target: std::cmp::Ordering) {
    if matches!(candidate, Value::Null) {
        return;
    }
    let Some(cand_key) = reddb_types::value_to_canonical_key(candidate) else {
        return;
    };
    match current {
        None => {
            *current = Some(candidate.clone());
        }
        Some(cur) => {
            let Some(cur_key) = reddb_types::value_to_canonical_key(cur) else {
                *current = Some(candidate.clone());
                return;
            };
            // Only compare within the same canonical family — cross-family
            // ordering would silently coerce, masking shape bugs.
            if cur_key.family() != cand_key.family() {
                return;
            }
            if cand_key.cmp(&cur_key) == target {
                *current = Some(candidate.clone());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::super::aggregate_value_to_f64;
    use reddb_types::Value;

    /// The planner's `SUM`/`AVG` cast and the legacy path's must agree on
    /// every value they both accept.
    ///
    /// They did not: the planner cast `Decimal` raw while the legacy path
    /// divided by the fixed scale, so the same `SUM` over a `DECIMAL` column
    /// returned values 10^4 apart depending on the route chosen (#2058). Both
    /// now delegate to `schema::decimal_to_f64`; this pins them so a future
    /// edit to one match arm cannot silently reintroduce the split.
    #[test]
    fn decimal_sum_agrees_with_the_legacy_aggregate_path() {
        let cases = [
            (Value::Decimal(387_600), Some(38.76)),
            (Value::Decimal(-771_500), Some(-77.15)),
            (Value::Decimal(1_234_567), Some(123.4567)),
            (Value::Decimal(0), Some(0.0)),
            (
                Value::Decimal(i64::MAX),
                Some(reddb_types::types::decimal_to_f64(i64::MAX)),
            ),
            (
                Value::Decimal(i64::MIN),
                Some(reddb_types::types::decimal_to_f64(i64::MIN)),
            ),
            (Value::Integer(42), Some(42.0)),
            (Value::UnsignedInteger(7), Some(7.0)),
            (Value::Float(1.5), Some(1.5)),
            (Value::Null, None),
        ];
        for (value, expected) in cases {
            assert_eq!(
                aggregate_value_to_f64(&value),
                expected,
                "shared aggregate cast changed for {value:?}"
            );
        }
    }

    enum Expected {
        Number(Option<f64>),
        Nan,
    }

    /// The planner and legacy routes must accept the same complete value matrix.
    #[test]
    fn planner_and_legacy_agree_on_the_complete_value_matrix() {
        let cases = [
            (Value::Integer(-3), Expected::Number(Some(-3.0))),
            (Value::UnsignedInteger(4), Expected::Number(Some(4.0))),
            (Value::BigInt(9), Expected::Number(Some(9.0))),
            (Value::Float(1.5), Expected::Number(Some(1.5))),
            (Value::Float(f64::NAN), Expected::Nan),
            (
                Value::Float(f64::INFINITY),
                Expected::Number(Some(f64::INFINITY)),
            ),
            (
                Value::Float(f64::NEG_INFINITY),
                Expected::Number(Some(f64::NEG_INFINITY)),
            ),
            (Value::Decimal(1_234_567), Expected::Number(Some(123.4567))),
            (Value::Boolean(true), Expected::Number(Some(1.0))),
            (Value::Boolean(false), Expected::Number(Some(0.0))),
            (Value::Text("not numeric".into()), Expected::Number(None)),
            (Value::Null, Expected::Number(None)),
        ];

        for (value, expected) in cases {
            let converted = aggregate_value_to_f64(&value);
            match expected {
                Expected::Number(expected) => {
                    assert_eq!(converted, expected, "unexpected cast for {value:?}")
                }
                Expected::Nan => assert!(converted.is_some_and(f64::is_nan)),
            }
        }
    }

    /// A `DECIMAL` read for aggregation must equal what the engine renders and
    /// compares it as — scale 4, the value the write side parses at.
    #[test]
    fn decimal_aggregation_matches_the_rendered_value() {
        // `Decimal(1_234_567)` renders as "123.4567" (pinned in entity_json).
        assert_eq!(
            aggregate_value_to_f64(&Value::Decimal(1_234_567)),
            Some(123.4567)
        );
    }
}