vibesql-executor 0.2.0

//! SIMD-accelerated aggregate computation via Arrow
//!
//! This module provides SIMD-accelerated evaluation for expression aggregates
//! using Apache Arrow. It converts rows to RecordBatch, evaluates expressions
//! using SIMD operations, and aggregates results using Arrow compute kernels.
//!
//! The SIMD path provides 4-8x performance improvement for large datasets
//! (>= 100 rows) compared to row-by-row evaluation.

use vibesql_ast::Expression;
use vibesql_storage::Row;
use vibesql_types::SqlValue;

use super::super::AggregateOp;
use crate::{errors::ExecutorError, schema::CombinedSchema};

/// Try to compute aggregate using SIMD-accelerated evaluation
///
/// This function converts rows to Arrow RecordBatch, evaluates the expression
/// using SIMD operations, and aggregates the result.
///
/// Returns Ok(value) if SIMD path succeeds, Err(_) if it fails (caller falls back to scalar).
pub fn try_simd_aggregate(
    rows: &[Row],
    expr: &Expression,
    op: AggregateOp,
    schema: &CombinedSchema,
) -> Result<SqlValue, ExecutorError> {
    log::debug!("[SIMD-Arrow] try_simd_aggregate called: {} rows, op={:?}", rows.len(), op);
    use crate::select::vectorized::{evaluate_arithmetic_simd, rows_to_record_batch};

    // Extract column names from schema (in order)
    let mut column_names = vec![String::new(); schema.total_columns];
    for (start_idx, table_schema) in schema.table_schemas.values() {
        for (col_idx, col) in table_schema.columns.iter().enumerate() {
            column_names[start_idx + col_idx] = col.name.clone();
        }
    }

    // Convert rows to RecordBatch
    let batch = rows_to_record_batch(rows, &column_names).map_err(|e| {
        ExecutorError::ArrowDowncastError {
            expected_type: "RecordBatch".to_string(),
            context: format!("rows_to_record_batch: {:?}", e),
        }
    })?;

    // Evaluate expression using SIMD
    let result_array = evaluate_arithmetic_simd(&batch, expr)?;

    // Aggregate the result array
    match op {
        AggregateOp::Sum => aggregate_sum(&result_array),
        AggregateOp::Count => aggregate_count(&result_array),
        AggregateOp::Avg => aggregate_avg(rows, expr, schema),
        AggregateOp::Min => aggregate_min(&result_array),
        AggregateOp::Max => aggregate_max(&result_array),
    }
}

/// Aggregate SUM using Arrow compute kernels
fn aggregate_sum(result_array: &arrow::array::ArrayRef) -> Result<SqlValue, ExecutorError> {
    use arrow::{
        array::{Float64Array, Int64Array},
        compute::sum,
    };

    match result_array.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let arr = result_array.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "SIMD aggregate".to_string(),
                }
            })?;
            let sum_val = sum(arr).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "SUM".to_string(),
                reason: "returned None".to_string(),
            })?;
            // SQLite's SUM() preserves integer type for integer inputs
            Ok(SqlValue::Integer(sum_val))
        }
        arrow::datatypes::DataType::Float64 => {
            let arr = result_array.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "SIMD aggregate".to_string(),
                }
            })?;
            let sum_val = sum(arr).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "SUM".to_string(),
                reason: "returned None".to_string(),
            })?;
            // SQLite's SUM() always returns REAL (Numeric)
            Ok(SqlValue::Numeric(sum_val))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "SUM".to_string(),
            array_type: format!("{:?}", result_array.data_type()),
        }),
    }
}

/// Aggregate COUNT (non-null values)
fn aggregate_count(result_array: &arrow::array::ArrayRef) -> Result<SqlValue, ExecutorError> {
    let non_null_count = result_array.len() - result_array.null_count();
    Ok(SqlValue::Integer(non_null_count as i64))
}

/// Aggregate AVG = SUM / COUNT
fn aggregate_avg(
    rows: &[Row],
    expr: &Expression,
    schema: &CombinedSchema,
) -> Result<SqlValue, ExecutorError> {
    let sum_result = try_simd_aggregate(rows, expr, AggregateOp::Sum, schema)?;
    let count_result = try_simd_aggregate(rows, expr, AggregateOp::Count, schema)?;

    match (sum_result, count_result) {
        (SqlValue::Numeric(sum), SqlValue::Integer(count)) if count > 0 => {
            Ok(SqlValue::Double(sum / count as f64))
        }
        (SqlValue::Double(sum), SqlValue::Integer(count)) if count > 0 => {
            Ok(SqlValue::Double(sum / count as f64))
        }
        (SqlValue::Integer(sum), SqlValue::Integer(count)) if count > 0 => {
            Ok(SqlValue::Double(sum as f64 / count as f64))
        }
        _ => Ok(SqlValue::Null),
    }
}

/// Aggregate MIN using Arrow compute kernels
fn aggregate_min(result_array: &arrow::array::ArrayRef) -> Result<SqlValue, ExecutorError> {
    use arrow::{
        array::{Float64Array, Int64Array},
        compute::min,
    };

    match result_array.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let arr = result_array.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "SIMD aggregate".to_string(),
                }
            })?;
            let min_val = min(arr).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MIN".to_string(),
                reason: "returned None".to_string(),
            })?;
            Ok(SqlValue::Integer(min_val))
        }
        arrow::datatypes::DataType::Float64 => {
            let arr = result_array.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "SIMD aggregate".to_string(),
                }
            })?;
            let min_val = min(arr).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MIN".to_string(),
                reason: "returned None".to_string(),
            })?;
            Ok(SqlValue::Double(min_val))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "MIN".to_string(),
            array_type: format!("{:?}", result_array.data_type()),
        }),
    }
}

/// Aggregate MAX using Arrow compute kernels
fn aggregate_max(result_array: &arrow::array::ArrayRef) -> Result<SqlValue, ExecutorError> {
    use arrow::{
        array::{Float64Array, Int64Array},
        compute::max,
    };

    match result_array.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let arr = result_array.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "SIMD aggregate".to_string(),
                }
            })?;
            let max_val = max(arr).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MAX".to_string(),
                reason: "returned None".to_string(),
            })?;
            Ok(SqlValue::Integer(max_val))
        }
        arrow::datatypes::DataType::Float64 => {
            let arr = result_array.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "SIMD aggregate".to_string(),
                }
            })?;
            let max_val = max(arr).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MAX".to_string(),
                reason: "returned None".to_string(),
            })?;
            Ok(SqlValue::Double(max_val))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "MAX".to_string(),
            array_type: format!("{:?}", result_array.data_type()),
        }),
    }
}