AkurAI-Framework / crates/sqlexec/src/exec.rs

Latest change fa37e2544d31a1199954a395f9e5b3ebb237b17f - feat(sql): aggregates (COUNT/SUM/AVG/MIN/MAX) + GROUP BY (gpt-5.5) by Ólafur Búi Ólafsson

//! The top-level [`Executor`]: owns a [`BTree`] and runs SQL text against it.

use std::path::Path;

use akurai_sql::{
    parse, AggregateArg, AggregateExpr, AggregateFunction, CreateTableStmt, DataType, DeleteStmt,
    Direction, InsertStmt, Literal, SelectColumn, SelectStmt, Statement, UpdateStmt,
};
use akurai_storage::BTree;

use crate::catalog::Schema;
use crate::error::ExecError;
use crate::row::{decode_row, encode_row};
use crate::value::Value;
use crate::{eval, ResultSet};

/// What a successful statement produced.
#[derive(Debug, PartialEq)]
pub enum Outcome {
    /// A DDL statement (`CREATE TABLE`) succeeded.
    Created(String),
    /// A DML statement wrote rows; carries the count affected.
    Affected(usize),
    /// A query produced a result set.
    Rows(ResultSet),
}

/// A SQL engine over a single B+tree file.
pub struct Executor {
    tree: BTree,
}

impl Executor {
    /// Open (or create) the database file backing this engine.
    pub fn open(path: impl AsRef<Path>) -> Result<Executor, ExecError> {
        Ok(Executor {
            tree: BTree::open(path)?,
        })
    }

    /// Parse and run one SQL statement.
    pub fn execute(&mut self, sql: &str) -> Result<Outcome, ExecError> {
        match parse(sql)? {
            Statement::CreateTable(stmt) => self.create_table(stmt),
            Statement::Insert(stmt) => self.insert(stmt),
            Statement::Update(stmt) => self.update(stmt),
            Statement::Delete(stmt) => self.delete(stmt),
            Statement::Select(stmt) => self.select(stmt),
        }
    }

    // ---- DDL --------------------------------------------------------------

    fn create_table(&mut self, stmt: CreateTableStmt) -> Result<Outcome, ExecError> {
        let key = Schema::catalog_key(&stmt.table);
        if self.tree.get(&key)?.is_some() {
            return Err(ExecError::TableExists(stmt.table));
        }
        let schema = Schema::from_defs(&stmt.table, &stmt.columns);
        self.tree.insert(&key, &schema.encode())?;
        // Seed the rowid counter at 0 so INSERT can read it unconditionally.
        self.tree
            .insert(&Schema::seq_key(&stmt.table), &0u64.to_be_bytes())?;
        self.tree.commit()?;
        Ok(Outcome::Created(stmt.table))
    }

    // ---- INSERT -----------------------------------------------------------

    fn insert(&mut self, stmt: InsertStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        if stmt.columns.len() != stmt.values.len() {
            return Err(ExecError::ColumnCountMismatch {
                expected: stmt.columns.len(),
                got: stmt.values.len(),
            });
        }

        // Place each supplied value into its schema slot; unmentioned columns
        // default to NULL.
        let mut cells: Vec<Value> = vec![Value::Null; schema.columns.len()];
        for (name, lit) in stmt.columns.iter().zip(stmt.values.iter()) {
            let idx = schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone()))?;
            cells[idx] = coerce(lit, &schema.columns[idx])?;
        }

        // Enforce NOT NULL after placement.
        for (i, col) in schema.columns.iter().enumerate() {
            if col.not_null && cells[i].is_null() {
                return Err(ExecError::NotNull(col.name.clone()));
            }
        }

        let rowid = self.next_rowid(&stmt.table)?;
        let key = data_key(&stmt.table, rowid);
        self.tree.insert(&key, &encode_row(&cells))?;
        self.tree.commit()?;
        Ok(Outcome::Affected(1))
    }

    // ---- UPDATE -----------------------------------------------------------

    fn update(&mut self, stmt: UpdateStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        // Resolve every assignment up front: validates each target column exists
        // and that its value coerces to the declared type before touching rows.
        let mut sets: Vec<(usize, Value)> = Vec::with_capacity(stmt.assignments.len());
        for (name, lit) in &stmt.assignments {
            let idx = schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone()))?;
            let value = coerce(lit, &schema.columns[idx])?;
            sets.push((idx, value));
        }

        // Scan, filter by WHERE, and stage the rewritten rows (same key/rowid).
        let (start, end) = data_bounds(&stmt.table);
        let mut rewrites: Vec<(Vec<u8>, Vec<Value>)> = Vec::new();
        for (key, bytes) in self.tree.range(&start, &end)? {
            let mut row = decode_row(&bytes)?;
            if !matches_where(&stmt.where_clause, &schema, &row)? {
                continue;
            }
            for (idx, value) in &sets {
                row[*idx] = value.clone();
            }
            // Re-enforce NOT NULL after applying the assignments.
            for (i, col) in schema.columns.iter().enumerate() {
                if col.not_null && row[i].is_null() {
                    return Err(ExecError::NotNull(col.name.clone()));
                }
            }
            rewrites.push((key, row));
        }

        let n = rewrites.len();
        for (key, row) in rewrites {
            self.tree.insert(&key, &encode_row(&row))?;
        }
        self.tree.commit()?;
        Ok(Outcome::Affected(n))
    }

    // ---- DELETE -----------------------------------------------------------

    fn delete(&mut self, stmt: DeleteStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        let (start, end) = data_bounds(&stmt.table);
        let mut victims: Vec<Vec<u8>> = Vec::new();
        for (key, bytes) in self.tree.range(&start, &end)? {
            let row = decode_row(&bytes)?;
            if matches_where(&stmt.where_clause, &schema, &row)? {
                victims.push(key);
            }
        }

        let n = victims.len();
        for key in victims {
            self.tree.delete(&key)?;
        }
        self.tree.commit()?;
        Ok(Outcome::Affected(n))
    }

    // ---- SELECT -----------------------------------------------------------

    fn select(&mut self, stmt: SelectStmt) -> Result<Outcome, ExecError> {
        let schema = self.load_schema(&stmt.table)?;

        // 1. Scan the table's data range and decode rows.
        let (start, end) = data_bounds(&stmt.table);
        let mut rows: Vec<Vec<Value>> = Vec::new();
        for (_k, v) in self.tree.range(&start, &end)? {
            rows.push(decode_row(&v)?);
        }

        // 2. WHERE filter.
        if let Some(pred) = &stmt.where_clause {
            let mut kept = Vec::with_capacity(rows.len());
            for row in rows {
                if eval::eval_predicate(pred, &schema, &row)? {
                    kept.push(row);
                }
            }
            rows = kept;
        }

        if is_aggregate_query(&stmt) {
            return aggregate_select(&stmt, &schema, rows).map(Outcome::Rows);
        }

        // 3. ORDER BY.
        if let Some(order) = &stmt.order_by {
            let idx = schema
                .index_of(&order.column)
                .ok_or_else(|| ExecError::UnknownColumn(order.column.clone()))?;
            rows.sort_by(|a, b| {
                let ord = a[idx].order_key(&b[idx]);
                match order.direction {
                    Direction::Asc => ord,
                    Direction::Desc => ord.reverse(),
                }
            });
        }

        // 4. LIMIT.
        if let Some(limit) = stmt.limit {
            let n = limit.max(0) as usize;
            rows.truncate(n);
        }

        // 5. Projection.
        let proj = projection(&stmt.columns, &schema)?;
        let columns: Vec<String> = proj
            .iter()
            .map(|i| schema.columns[*i].name.clone())
            .collect();
        let out_rows: Vec<Vec<Value>> = rows
            .into_iter()
            .map(|row| proj.iter().map(|i| row[*i].clone()).collect())
            .collect();

        Ok(Outcome::Rows(ResultSet {
            columns,
            rows: out_rows,
        }))
    }

    // ---- helpers ----------------------------------------------------------

    fn load_schema(&mut self, table: &str) -> Result<Schema, ExecError> {
        let bytes = self
            .tree
            .get(&Schema::catalog_key(table))?
            .ok_or_else(|| ExecError::NoSuchTable(table.to_string()))?;
        Schema::decode(table, &bytes)
    }

    /// Read, bump, and persist the per-table rowid counter; returns the rowid
    /// to use for the new row (1-based).
    fn next_rowid(&mut self, table: &str) -> Result<u64, ExecError> {
        let key = Schema::seq_key(table);
        let current = match self.tree.get(&key)? {
            Some(b) if b.len() == 8 => {
                let mut a = [0u8; 8];
                a.copy_from_slice(&b);
                u64::from_be_bytes(a)
            }
            Some(_) => return Err(ExecError::Corrupt("rowid counter wrong size")),
            None => 0,
        };
        let next = current + 1;
        self.tree.insert(&key, &next.to_be_bytes())?;
        Ok(next)
    }
}

/// Whether `row` satisfies an optional `WHERE` predicate. A missing predicate
/// matches every row (the bare `UPDATE`/`DELETE` whole-table case).
fn matches_where(
    where_clause: &Option<akurai_sql::Expr>,
    schema: &Schema,
    row: &[Value],
) -> Result<bool, ExecError> {
    match where_clause {
        Some(pred) => eval::eval_predicate(pred, schema, row),
        None => Ok(true),
    }
}

/// Coerce a parser literal into the column's declared type, or reject it.
fn coerce(lit: &Literal, col: &crate::catalog::Column) -> Result<Value, ExecError> {
    match (lit, col.ty) {
        (Literal::Null, _) => Ok(Value::Null),
        (Literal::Int(i), DataType::Integer) => Ok(Value::Int(*i)),
        // Widen an integer literal into a REAL column.
        (Literal::Int(i), DataType::Real) => Ok(Value::Float(*i as f64)),
        (Literal::Float(f), DataType::Real) => Ok(Value::Float(*f)),
        (Literal::Str(s), DataType::Text) => Ok(Value::Str(s.clone())),
        _ => Err(ExecError::TypeMismatch {
            column: col.name.clone(),
        }),
    }
}

/// Resolve the projection list into schema column indices.
fn projection(cols: &[SelectColumn], schema: &Schema) -> Result<Vec<usize>, ExecError> {
    // `SELECT *` (or a list containing `*`) projects every column in order.
    if cols.iter().any(|c| matches!(c, SelectColumn::All)) {
        return Ok((0..schema.columns.len()).collect());
    }
    cols.iter()
        .map(|c| match c {
            SelectColumn::All => unreachable!(),
            SelectColumn::Named(name) => schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone())),
            SelectColumn::Aggregate(_) => Err(ExecError::Unsupported(
                "aggregate projection in non-aggregate SELECT",
            )),
        })
        .collect()
}

fn is_aggregate_query(stmt: &SelectStmt) -> bool {
    stmt.group_by.is_some()
        || stmt
            .columns
            .iter()
            .any(|c| matches!(c, SelectColumn::Aggregate(_)))
}

fn aggregate_select(
    stmt: &SelectStmt,
    schema: &Schema,
    rows: Vec<Vec<Value>>,
) -> Result<ResultSet, ExecError> {
    let plan = AggregatePlan::new(stmt, schema)?;
    let mut groups = build_groups(&plan, rows)?;

    if let Some(order) = &stmt.order_by {
        if plan.group_name.as_deref() == Some(order.column.as_str()) {
            groups.sort_by(|a, b| {
                let ord = match (&a.key, &b.key) {
                    (Some(a), Some(b)) => a.order_key(b),
                    _ => std::cmp::Ordering::Equal,
                };
                match order.direction {
                    Direction::Asc => ord,
                    Direction::Desc => ord.reverse(),
                }
            });
        }
    }

    let mut out_rows: Vec<Vec<Value>> = groups
        .iter()
        .map(|group| plan.materialize_row(group))
        .collect::<Result<_, _>>()?;
    let columns = plan.output_columns();

    if let Some(order) = &stmt.order_by {
        if plan.group_name.as_deref() != Some(order.column.as_str()) {
            if let Some(idx) = columns.iter().position(|c| c == &order.column) {
                out_rows.sort_by(|a, b| {
                    let ord = a[idx].order_key(&b[idx]);
                    match order.direction {
                        Direction::Asc => ord,
                        Direction::Desc => ord.reverse(),
                    }
                });
            } else if schema.index_of(&order.column).is_none() {
                return Err(ExecError::UnknownColumn(order.column.clone()));
            } else if stmt.group_by.is_some() {
                return Err(ExecError::Unsupported(
                    "ORDER BY in grouped SELECT must reference the group column or output",
                ));
            }
        }
    }

    if let Some(limit) = stmt.limit {
        let n = limit.max(0) as usize;
        out_rows.truncate(n);
    }

    Ok(ResultSet {
        columns,
        rows: out_rows,
    })
}

struct AggregatePlan {
    group_name: Option<String>,
    group_idx: Option<usize>,
    outputs: Vec<AggregateOutput>,
    aggregates: Vec<PlannedAggregate>,
}

enum AggregateOutput {
    GroupColumn(String),
    Aggregate(usize),
}

struct PlannedAggregate {
    expr: AggregateExpr,
    arg_idx: Option<usize>,
}

impl AggregatePlan {
    fn new(stmt: &SelectStmt, schema: &Schema) -> Result<AggregatePlan, ExecError> {
        let group_idx = match &stmt.group_by {
            Some(name) => Some(
                schema
                    .index_of(name)
                    .ok_or_else(|| ExecError::UnknownColumn(name.clone()))?,
            ),
            None => None,
        };

        let mut outputs = Vec::with_capacity(stmt.columns.len());
        let mut aggregates = Vec::new();
        for column in &stmt.columns {
            match column {
                SelectColumn::All => {
                    return Err(ExecError::Unsupported(
                        "SELECT * cannot be mixed with aggregates or GROUP BY",
                    ))
                }
                SelectColumn::Named(name) => {
                    if stmt.group_by.as_ref() == Some(name) {
                        outputs.push(AggregateOutput::GroupColumn(name.clone()));
                    } else if stmt.group_by.is_some() {
                        return Err(ExecError::Unsupported(
                            "grouped SELECT can only project the group column and aggregates",
                        ));
                    } else {
                        return Err(ExecError::Unsupported(
                            "aggregate SELECT can only project aggregates without GROUP BY",
                        ));
                    }
                }
                SelectColumn::Aggregate(expr) => {
                    let idx = aggregates.len();
                    aggregates.push(plan_aggregate(expr, schema)?);
                    outputs.push(AggregateOutput::Aggregate(idx));
                }
            }
        }

        Ok(AggregatePlan {
            group_name: stmt.group_by.clone(),
            group_idx,
            outputs,
            aggregates,
        })
    }

    fn output_columns(&self) -> Vec<String> {
        self.outputs
            .iter()
            .map(|output| match output {
                AggregateOutput::GroupColumn(name) => name.clone(),
                AggregateOutput::Aggregate(idx) => aggregate_label(&self.aggregates[*idx].expr),
            })
            .collect()
    }

    fn materialize_row(&self, group: &GroupState) -> Result<Vec<Value>, ExecError> {
        self.outputs
            .iter()
            .map(|output| match output {
                AggregateOutput::GroupColumn(_) => Ok(group.key.clone().unwrap_or(Value::Null)),
                AggregateOutput::Aggregate(idx) => Ok(group.states[*idx].finish()),
            })
            .collect()
    }
}

fn plan_aggregate(expr: &AggregateExpr, schema: &Schema) -> Result<PlannedAggregate, ExecError> {
    let arg_idx = match &expr.arg {
        AggregateArg::All => {
            if expr.func != AggregateFunction::Count {
                return Err(ExecError::Unsupported("only COUNT accepts '*'"));
            }
            None
        }
        AggregateArg::Column(name) => {
            let idx = schema
                .index_of(name)
                .ok_or_else(|| ExecError::UnknownColumn(name.clone()))?;
            if matches!(expr.func, AggregateFunction::Sum | AggregateFunction::Avg)
                && schema.columns[idx].ty == DataType::Text
            {
                return Err(ExecError::TypeMismatch {
                    column: name.clone(),
                });
            }
            Some(idx)
        }
    };
    Ok(PlannedAggregate {
        expr: expr.clone(),
        arg_idx,
    })
}

struct GroupState {
    key: Option<Value>,
    states: Vec<AggregateState>,
}

fn build_groups(plan: &AggregatePlan, rows: Vec<Vec<Value>>) -> Result<Vec<GroupState>, ExecError> {
    let mut groups: Vec<GroupState> = Vec::new();

    if let Some(group_idx) = plan.group_idx {
        for row in rows {
            let key = row.get(group_idx).cloned().unwrap_or(Value::Null);
            let idx = match groups
                .iter()
                .position(|group| group.key.as_ref() == Some(&key))
            {
                Some(idx) => idx,
                None => {
                    groups.push(GroupState {
                        key: Some(key),
                        states: new_aggregate_states(plan),
                    });
                    groups.len() - 1
                }
            };
            update_group(plan, &mut groups[idx], &row)?;
        }
    } else {
        groups.push(GroupState {
            key: None,
            states: new_aggregate_states(plan),
        });
        for row in rows {
            update_group(plan, &mut groups[0], &row)?;
        }
    }

    Ok(groups)
}

fn new_aggregate_states(plan: &AggregatePlan) -> Vec<AggregateState> {
    plan.aggregates
        .iter()
        .map(|agg| AggregateState::new(agg.expr.func))
        .collect()
}

fn update_group(
    plan: &AggregatePlan,
    group: &mut GroupState,
    row: &[Value],
) -> Result<(), ExecError> {
    for (state, agg) in group.states.iter_mut().zip(plan.aggregates.iter()) {
        let value = agg.arg_idx.and_then(|idx| row.get(idx));
        state.update(value)?;
    }
    Ok(())
}

#[derive(Clone)]
struct AggregateState {
    func: AggregateFunction,
    count: i64,
    numeric_count: i64,
    int_sum: i64,
    float_sum: f64,
    saw_float: bool,
    int_overflow: bool,
    min: Option<Value>,
    max: Option<Value>,
}

impl AggregateState {
    fn new(func: AggregateFunction) -> AggregateState {
        AggregateState {
            func,
            count: 0,
            numeric_count: 0,
            int_sum: 0,
            float_sum: 0.0,
            saw_float: false,
            int_overflow: false,
            min: None,
            max: None,
        }
    }

    fn update(&mut self, value: Option<&Value>) -> Result<(), ExecError> {
        match self.func {
            AggregateFunction::Count => {
                if match value {
                    Some(v) => !v.is_null(),
                    None => true,
                } {
                    self.count = self.count.saturating_add(1);
                }
            }
            AggregateFunction::Sum | AggregateFunction::Avg => {
                let Some(value) = value else {
                    return Ok(());
                };
                match value {
                    Value::Null => {}
                    Value::Int(i) => self.add_int(*i),
                    Value::Float(f) => self.add_float(*f),
                    Value::Str(_) => {
                        return Err(ExecError::TypeMismatch {
                            column: aggregate_label_func(self.func).to_string(),
                        })
                    }
                }
            }
            AggregateFunction::Min => {
                if let Some(value) = value.filter(|v| !v.is_null()) {
                    let replace = match &self.min {
                        Some(current) => value.order_key(current).is_lt(),
                        None => true,
                    };
                    if replace {
                        self.min = Some(value.clone());
                    }
                }
            }
            AggregateFunction::Max => {
                if let Some(value) = value.filter(|v| !v.is_null()) {
                    let replace = match &self.max {
                        Some(current) => value.order_key(current).is_gt(),
                        None => true,
                    };
                    if replace {
                        self.max = Some(value.clone());
                    }
                }
            }
        }
        Ok(())
    }

    fn add_int(&mut self, value: i64) {
        self.numeric_count = self.numeric_count.saturating_add(1);
        self.float_sum += value as f64;
        match self.int_sum.checked_add(value) {
            Some(sum) => self.int_sum = sum,
            None => self.int_overflow = true,
        }
    }

    fn add_float(&mut self, value: f64) {
        self.numeric_count = self.numeric_count.saturating_add(1);
        self.float_sum += value;
        self.saw_float = true;
    }

    fn finish(&self) -> Value {
        match self.func {
            AggregateFunction::Count => Value::Int(self.count),
            AggregateFunction::Sum => {
                if self.numeric_count == 0 {
                    Value::Null
                } else if self.saw_float || self.int_overflow {
                    Value::Float(self.float_sum)
                } else {
                    Value::Int(self.int_sum)
                }
            }
            AggregateFunction::Avg => {
                if self.numeric_count == 0 {
                    Value::Null
                } else {
                    Value::Float(self.float_sum / self.numeric_count as f64)
                }
            }
            AggregateFunction::Min => self.min.clone().unwrap_or(Value::Null),
            AggregateFunction::Max => self.max.clone().unwrap_or(Value::Null),
        }
    }
}

fn aggregate_label(expr: &AggregateExpr) -> String {
    let arg = match &expr.arg {
        AggregateArg::All => "*".to_string(),
        AggregateArg::Column(name) => name.clone(),
    };
    format!("{}({arg})", aggregate_label_func(expr.func))
}

fn aggregate_label_func(func: AggregateFunction) -> &'static str {
    match func {
        AggregateFunction::Count => "COUNT",
        AggregateFunction::Sum => "SUM",
        AggregateFunction::Avg => "AVG",
        AggregateFunction::Min => "MIN",
        AggregateFunction::Max => "MAX",
    }
}

/// The data key for `table`/`rowid`: `<table>:<rowid_be>`.
fn data_key(table: &str, rowid: u64) -> Vec<u8> {
    let mut k = Vec::with_capacity(table.len() + 9);
    k.extend_from_slice(table.as_bytes());
    k.push(b':');
    k.extend_from_slice(&rowid.to_be_bytes());
    k
}

/// The half-open `[start, end)` key window covering every row of `table`.
fn data_bounds(table: &str) -> (Vec<u8>, Vec<u8>) {
    let mut prefix = Vec::with_capacity(table.len() + 1);
    prefix.extend_from_slice(table.as_bytes());
    prefix.push(b':');
    let end = prefix_upper(&prefix);
    (prefix, end)
}

/// The smallest key strictly greater than every key beginning with `prefix`.
fn prefix_upper(prefix: &[u8]) -> Vec<u8> {
    let mut end = prefix.to_vec();
    while let Some(last) = end.last_mut() {
        if *last != 0xFF {
            *last += 1;
            return end;
        }
        end.pop();
    }
    // All-0xFF prefix: no upper bound shorter than infinity. Fall back to a
    // sentinel that the storage layer treats as "scan to the end".
    vec![0xFF; prefix.len() + 1]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn prefix_upper_increments_last_byte() {
        assert_eq!(prefix_upper(b"users:"), b"users;".to_vec());
    }

    #[test]
    fn data_key_orders_by_rowid() {
        assert!(data_key("t", 1) < data_key("t", 2));
        assert!(data_key("t", 2) < data_key("t", 256));
    }
}