tpt-archon-relational 0.1.0

//! `GROUP BY` + aggregate function evaluation.

use alloc::string::{String, ToString};
use alloc::vec::Vec;

use crate::parser::AggregateFunc;

use super::value::{ExecError, ResultSet, Row, Value};

/// Resolves an [`Expr::Agg`](crate::parser::Expr::Agg) to its default column
/// name (the alias generated by the parser when no explicit `AS` is given).
/// For explicit aliases, the user references the alias directly in `HAVING`
/// (e.g. `HAVING cnt > 1`).
pub(crate) fn agg_default_alias(func: AggregateFunc, column: &str) -> String {
    let tag = match func {
        AggregateFunc::Count => "count",
        AggregateFunc::Sum => "sum",
        AggregateFunc::Avg => "avg",
        AggregateFunc::Min => "min",
        AggregateFunc::Max => "max",
    };
    if column == "*" {
        tag.to_string()
    } else {
        alloc::format!("{tag}({column})")
    }
}

/// Runs `GROUP BY` + aggregate projection over an in-memory `(columns, rows)`
/// table, producing the aggregated result. Shared by the executor's own
/// [`PlanNode::Aggregate`](crate::planner::PlanNode::Aggregate) and
/// [`Database::run_select`](crate::database::Database::run_select)'s
/// pre-planner aggregate step so the two paths don't duplicate this logic.
pub fn aggregate_table(
    columns: &[String],
    rows: &[Row],
    group_by: &[String],
    aggregates: &[(String, AggregateFunc, String)],
) -> Result<ResultSet, ExecError> {
    if group_by.is_empty() {
        let mut out_row = Vec::new();
        let mut out_cols = Vec::new();
        for (alias, func, col) in aggregates {
            out_row.push(eval_aggregate(*func, col, columns, rows)?);
            out_cols.push(alias.clone());
        }
        return Ok(ResultSet {
            columns: out_cols,
            rows: alloc::vec![out_row],
            ..Default::default()
        });
    }

    let gb_indices: Vec<usize> = group_by
        .iter()
        .map(|c| {
            columns
                .iter()
                .position(|ic| ic == c)
                .ok_or_else(|| ExecError::GroupByColumnNotFound(c.clone()))
        })
        .collect::<Result<_, _>>()?;

    let mut groups: alloc::collections::BTreeMap<Vec<Value>, Vec<Row>> =
        alloc::collections::BTreeMap::new();
    for row in rows {
        let key: Vec<Value> = gb_indices.iter().map(|&i| row[i].clone()).collect();
        groups.entry(key).or_default().push(row.clone());
    }

    let mut out_cols = group_by.to_vec();
    for (alias, _, _) in aggregates {
        out_cols.push(alias.clone());
    }
    let mut out_rows = Vec::new();
    for group_rows in groups.values() {
        let mut out_row = Vec::new();
        for &idx in &gb_indices {
            out_row.push(group_rows[0][idx].clone());
        }
        for (_alias, func, col) in aggregates {
            out_row.push(eval_aggregate(*func, col, columns, group_rows)?);
        }
        out_rows.push(out_row);
    }
    Ok(ResultSet {
        columns: out_cols,
        rows: out_rows,
        ..Default::default()
    })
}

/// Applies one aggregate function over `rows` (a full group, or — reused by
/// [`super::window`] — an arbitrary window frame's row slice).
pub(crate) fn eval_aggregate(
    func: AggregateFunc,
    column: &str,
    columns: &[String],
    rows: &[Row],
) -> Result<Value, ExecError> {
    if func == AggregateFunc::Count {
        return Ok(Value::Int(rows.len() as i64));
    }
    let idx = columns
        .iter()
        .position(|c| c == column)
        .ok_or_else(|| ExecError::UnknownColumn(column.to_string()))?;

    match func {
        AggregateFunc::Count => unreachable!(),
        AggregateFunc::Sum => {
            let sum: i64 = rows
                .iter()
                .filter_map(|r| match &r[idx] {
                    Value::Int(v) => Some(*v),
                    Value::Float(v) => Some(*v as i64),
                    _ => None,
                })
                .sum();
            Ok(Value::Int(sum))
        }
        AggregateFunc::Avg => {
            let vals: Vec<i64> = rows
                .iter()
                .filter_map(|r| match &r[idx] {
                    Value::Int(v) => Some(*v),
                    Value::Float(v) => Some(*v as i64),
                    _ => None,
                })
                .collect();
            if vals.is_empty() {
                Ok(Value::Int(0))
            } else {
                let avg = vals.iter().sum::<i64>() / vals.len() as i64;
                Ok(Value::Int(avg))
            }
        }
        AggregateFunc::Min => {
            let min = rows
                .iter()
                .filter_map(|r| match &r[idx] {
                    Value::Int(v) => Some(Value::Int(*v)),
                    Value::Float(v) => Some(Value::Float(*v)),
                    _ => None,
                })
                .min();
            Ok(min.unwrap_or(Value::Null))
        }
        AggregateFunc::Max => {
            let max = rows
                .iter()
                .filter_map(|r| match &r[idx] {
                    Value::Int(v) => Some(Value::Int(*v)),
                    Value::Float(v) => Some(Value::Float(*v)),
                    _ => None,
                })
                .max();
            Ok(max.unwrap_or(Value::Null))
        }
    }
}