vibesql-executor 0.2.0

//! Expression evaluation helpers for columnar batch execution
//!
//! This module provides functions for evaluating SQL expressions on columnar
//! data, including support for aggregate expressions like SUM(col_a * col_b).

use vibesql_ast::{BinaryOperator, Expression};
use vibesql_types::SqlValue;

use super::super::{
    aggregate::AggregateOp,
    batch::{ColumnArray, ColumnarBatch},
    simd_ops,
};
use crate::errors::ExecutorError;

/// Compute an aggregate over an expression on a ColumnarBatch
///
/// For simple binary operations like `col_a * col_b`, this evaluates the
/// expression element-wise on the columnar data and then aggregates.
pub fn compute_expression_aggregate_batch(
    batch: &ColumnarBatch,
    expr: &Expression,
    op: AggregateOp,
) -> Result<SqlValue, ExecutorError> {
    // Try to evaluate the expression as a simple binary operation on columns
    if let Expression::BinaryOp { left, op: bin_op, right } = expr {
        // For now, only support multiplication (most common in TPC-H)
        if *bin_op == BinaryOperator::Multiply {
            // Get column indices from left and right operands
            if let (Expression::ColumnRef(col_id1), Expression::ColumnRef(col_id2)) =
                (left.as_ref(), right.as_ref())
            {
                if col_id1.schema_canonical().is_none() && col_id2.schema_canonical().is_none() {
                    // Find column indices by name
                    let left_idx = batch.column_index_by_name(col_id1.column_canonical());
                    let right_idx = batch.column_index_by_name(col_id2.column_canonical());

                    if let (Some(l_idx), Some(r_idx)) = (left_idx, right_idx) {
                        return compute_multiply_aggregate(batch, l_idx, r_idx, op);
                    }
                }
            }
        }
    }

    // Fall back to row-by-row evaluation for complex expressions
    // This is slower but handles all cases
    let row_count = batch.row_count();
    let mut sum = 0.0f64;
    let mut count = 0i64;

    for row_idx in 0..row_count {
        // Evaluate expression for this row
        if let Ok(value) = eval_expr_on_batch(batch, expr, row_idx) {
            match value {
                SqlValue::Integer(v) => {
                    sum += v as f64;
                    count += 1;
                }
                SqlValue::Double(v) => {
                    sum += v;
                    count += 1;
                }
                SqlValue::Float(v) => {
                    sum += v as f64;
                    count += 1;
                }
                SqlValue::Bigint(v) => {
                    sum += v as f64;
                    count += 1;
                }
                SqlValue::Numeric(v) => {
                    sum += v;
                    count += 1;
                }
                SqlValue::Null => {}
                _ => {}
            }
        }
    }

    match op {
        AggregateOp::Sum => Ok(if count > 0 { SqlValue::Double(sum) } else { SqlValue::Null }),
        AggregateOp::Count => Ok(SqlValue::Integer(count)),
        AggregateOp::Avg => {
            Ok(if count > 0 { SqlValue::Double(sum / count as f64) } else { SqlValue::Null })
        }
        _ => Err(ExecutorError::UnsupportedExpression(
            "MIN/MAX not supported for expression aggregates".to_string(),
        )),
    }
}

/// Optimized path for SUM(col_a * col_b) - the most common expression aggregate in TPC-H
///
/// Uses SIMD-accelerated sum_product_f64 for maximum performance. This provides
/// ~4x speedup compared to the previous Vec<Option<f64>> approach.
pub fn compute_multiply_aggregate(
    batch: &ColumnarBatch,
    left_idx: usize,
    right_idx: usize,
    op: AggregateOp,
) -> Result<SqlValue, ExecutorError> {
    let left_col = batch.column(left_idx).ok_or_else(|| ExecutorError::ColumnarColumnNotFound {
        column_index: left_idx,
        batch_columns: batch.column_count(),
    })?;
    let right_col =
        batch.column(right_idx).ok_or_else(|| ExecutorError::ColumnarColumnNotFound {
            column_index: right_idx,
            batch_columns: batch.column_count(),
        })?;

    // Fast path: Both columns are Float64 - use SIMD-accelerated sum_product
    if let (
        ColumnArray::Float64(left_values, left_nulls),
        ColumnArray::Float64(right_values, right_nulls),
    ) = (left_col, right_col)
    {
        if left_values.len() != right_values.len() {
            return Err(ExecutorError::ColumnarLengthMismatch {
                context: "multiply_aggregate".to_string(),
                expected: left_values.len(),
                actual: right_values.len(),
            });
        }

        let (sum, count) = simd_ops::sum_product_f64_masked(
            left_values,
            right_values,
            left_nulls.as_ref().map(|v| v.as_slice()),
            right_nulls.as_ref().map(|v| v.as_slice()),
        );

        return match op {
            AggregateOp::Sum => Ok(if count > 0 { SqlValue::Double(sum) } else { SqlValue::Null }),
            AggregateOp::Count => Ok(SqlValue::Integer(count)),
            AggregateOp::Avg => {
                Ok(if count > 0 { SqlValue::Double(sum / count as f64) } else { SqlValue::Null })
            }
            _ => Err(ExecutorError::UnsupportedExpression(
                "MIN/MAX not supported for expression aggregates".to_string(),
            )),
        };
    }

    // Fast path: Both columns are Int64 - convert to f64 and use SIMD
    if let (
        ColumnArray::Int64(left_values, left_nulls),
        ColumnArray::Int64(right_values, right_nulls),
    ) = (left_col, right_col)
    {
        if left_values.len() != right_values.len() {
            return Err(ExecutorError::ColumnarLengthMismatch {
                context: "multiply_aggregate".to_string(),
                expected: left_values.len(),
                actual: right_values.len(),
            });
        }

        // Check if we have any nulls
        let has_left_nulls = left_nulls.as_ref().map_or(false, |n| n.iter().any(|&x| x));
        let has_right_nulls = right_nulls.as_ref().map_or(false, |n| n.iter().any(|&x| x));

        if !has_left_nulls && !has_right_nulls {
            // No nulls: use 4-accumulator SIMD pattern directly on i64
            let len = left_values.len();
            let (mut s0, mut s1, mut s2, mut s3) = (0i64, 0i64, 0i64, 0i64);
            let chunks = len / 4;

            for i in 0..chunks {
                let off = i * 4;
                s0 = s0.wrapping_add(left_values[off].wrapping_mul(right_values[off]));
                s1 = s1.wrapping_add(left_values[off + 1].wrapping_mul(right_values[off + 1]));
                s2 = s2.wrapping_add(left_values[off + 2].wrapping_mul(right_values[off + 2]));
                s3 = s3.wrapping_add(left_values[off + 3].wrapping_mul(right_values[off + 3]));
            }

            let mut sum = s0.wrapping_add(s1).wrapping_add(s2).wrapping_add(s3);
            for i in (chunks * 4)..len {
                sum = sum.wrapping_add(left_values[i].wrapping_mul(right_values[i]));
            }

            return match op {
                AggregateOp::Sum => {
                    // SQLite's SUM() preserves integer type for integer inputs
                    Ok(SqlValue::Integer(sum))
                }
                AggregateOp::Count => Ok(SqlValue::Integer(len as i64)),
                AggregateOp::Avg => Ok(if len > 0 {
                    SqlValue::Double(sum as f64 / len as f64)
                } else {
                    SqlValue::Null
                }),
                _ => Err(ExecutorError::UnsupportedExpression(
                    "MIN/MAX not supported for expression aggregates".to_string(),
                )),
            };
        }

        // Has nulls: use masked path
        let null_a = left_nulls.as_ref().map(|v| v.as_slice());
        let null_b = right_nulls.as_ref().map(|v| v.as_slice());

        let mut sum = 0i64;
        let mut count = 0i64;
        for i in 0..left_values.len() {
            let is_null_a = null_a.map_or(false, |n| n.get(i).copied().unwrap_or(false));
            let is_null_b = null_b.map_or(false, |n| n.get(i).copied().unwrap_or(false));
            if !is_null_a && !is_null_b {
                sum = sum.wrapping_add(left_values[i].wrapping_mul(right_values[i]));
                count += 1;
            }
        }

        return match op {
            AggregateOp::Sum => {
                // SQLite's SUM() preserves integer type for integer inputs
                Ok(if count > 0 { SqlValue::Integer(sum) } else { SqlValue::Null })
            }
            AggregateOp::Count => Ok(SqlValue::Integer(count)),
            AggregateOp::Avg => Ok(if count > 0 {
                SqlValue::Double(sum as f64 / count as f64)
            } else {
                SqlValue::Null
            }),
            _ => Err(ExecutorError::UnsupportedExpression(
                "MIN/MAX not supported for expression aggregates".to_string(),
            )),
        };
    }

    // Fallback: Mixed column types - use the slower Vec<Option<f64>> path
    let left_f64 = column_to_f64_vec(left_col)?;
    let right_f64 = column_to_f64_vec(right_col)?;

    if left_f64.len() != right_f64.len() {
        return Err(ExecutorError::ColumnarLengthMismatch {
            context: "multiply_aggregate".to_string(),
            expected: left_f64.len(),
            actual: right_f64.len(),
        });
    }

    // Vectorized multiply and aggregate
    let mut sum = 0.0f64;
    let mut count = 0i64;

    for i in 0..left_f64.len() {
        if let (Some(l), Some(r)) = (left_f64[i], right_f64[i]) {
            sum += l * r;
            count += 1;
        }
    }

    match op {
        AggregateOp::Sum => Ok(if count > 0 { SqlValue::Double(sum) } else { SqlValue::Null }),
        AggregateOp::Count => Ok(SqlValue::Integer(count)),
        AggregateOp::Avg => {
            Ok(if count > 0 { SqlValue::Double(sum / count as f64) } else { SqlValue::Null })
        }
        _ => Err(ExecutorError::UnsupportedExpression(
            "MIN/MAX not supported for expression aggregates".to_string(),
        )),
    }
}

/// Convert a ColumnArray to Vec<Option<f64>> for arithmetic operations
pub fn column_to_f64_vec(column: &ColumnArray) -> Result<Vec<Option<f64>>, ExecutorError> {
    match column {
        ColumnArray::Int64(values, nulls) => Ok(values
            .iter()
            .enumerate()
            .map(
                |(i, &v)| {
                    if nulls.as_ref().map_or(false, |n| n[i]) {
                        None
                    } else {
                        Some(v as f64)
                    }
                },
            )
            .collect()),
        ColumnArray::Float64(values, nulls) => Ok(values
            .iter()
            .enumerate()
            .map(|(i, &v)| if nulls.as_ref().map_or(false, |n| n[i]) { None } else { Some(v) })
            .collect()),
        ColumnArray::Mixed(values) => Ok(values
            .iter()
            .map(|v| match v {
                SqlValue::Integer(n) => Some(*n as f64),
                SqlValue::Bigint(n) => Some(*n as f64),
                SqlValue::Float(n) => Some(*n as f64),
                SqlValue::Double(n) => Some(*n),
                SqlValue::Numeric(n) => Some(*n),
                SqlValue::Null => None,
                _ => None,
            })
            .collect()),
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "column_to_f64".to_string(),
            array_type: format!("{:?}", std::mem::discriminant(column)),
        }),
    }
}

/// Evaluate an expression for a single row in a ColumnarBatch
pub fn eval_expr_on_batch(
    batch: &ColumnarBatch,
    expr: &Expression,
    row_idx: usize,
) -> Result<SqlValue, ExecutorError> {
    match expr {
        Expression::ColumnRef(col_id) => {
            if let Some(col_idx) = batch.column_index_by_name(col_id.column_canonical()) {
                batch.get_value(row_idx, col_idx)
            } else {
                Ok(SqlValue::Null)
            }
        }
        Expression::Literal(val) => Ok(val.clone()),
        Expression::BinaryOp { left, op, right } => {
            let left_val = eval_expr_on_batch(batch, left, row_idx)?;
            let right_val = eval_expr_on_batch(batch, right, row_idx)?;

            // Simple arithmetic evaluation
            match (left_val, right_val) {
                (SqlValue::Null, _) | (_, SqlValue::Null) => Ok(SqlValue::Null),
                (l, r) => {
                    let l_f64 = sql_value_to_f64(&l).ok_or_else(|| {
                        ExecutorError::ColumnarTypeMismatch {
                            operation: "binary_op".to_string(),
                            left_type: format!("{:?}", l),
                            right_type: None,
                        }
                    })?;
                    let r_f64 = sql_value_to_f64(&r).ok_or_else(|| {
                        ExecutorError::ColumnarTypeMismatch {
                            operation: "binary_op".to_string(),
                            left_type: format!("{:?}", r),
                            right_type: None,
                        }
                    })?;

                    let result = match op {
                        BinaryOperator::Plus => l_f64 + r_f64,
                        BinaryOperator::Minus => l_f64 - r_f64,
                        BinaryOperator::Multiply => l_f64 * r_f64,
                        BinaryOperator::Divide => l_f64 / r_f64,
                        _ => {
                            return Err(ExecutorError::UnsupportedExpression(format!(
                                "Unsupported binary operator: {:?}",
                                op
                            )))
                        }
                    };
                    Ok(SqlValue::Double(result))
                }
            }
        }
        _ => Err(ExecutorError::UnsupportedExpression(
            "Complex expression not supported".to_string(),
        )),
    }
}

/// Convert SqlValue to f64
pub fn sql_value_to_f64(val: &SqlValue) -> Option<f64> {
    match val {
        SqlValue::Integer(v) => Some(*v as f64),
        SqlValue::Bigint(v) => Some(*v as f64),
        SqlValue::Float(v) => Some(*v as f64),
        SqlValue::Double(v) => Some(*v),
        SqlValue::Numeric(v) => Some(*v),
        SqlValue::Smallint(v) => Some(*v as f64),
        _ => None,
    }
}