vibesql-executor 0.2.0

//! SIMD-accelerated aggregation using Arrow compute kernels
//!
//! This module provides vectorized aggregation operations using Apache Arrow's
//! compute kernels for SIMD acceleration.
//!
//! Note: This module is experimental/research code.

#![allow(dead_code)]

use arrow::{
    array::{Array, ArrayRef, Date32Array, Float64Array, Int64Array, TimestampMicrosecondArray},
    compute::kernels::aggregate::{max, min, sum},
    datatypes::TimeUnit,
    record_batch::RecordBatch,
};
use vibesql_types::SqlValue;

use crate::errors::ExecutorError;

/// Aggregate function types supported by SIMD execution
#[derive(Debug, Clone, PartialEq)]
pub enum AggregateFunction {
    Sum,
    Avg,
    Count,
    Min,
    Max,
}

/// Perform aggregation on a RecordBatch column using SIMD
///
/// This function uses Arrow's compute kernels which automatically leverage
/// SIMD instructions for aggregation operations.
pub fn aggregate_column_simd(
    batch: &RecordBatch,
    column_name: &str,
    function: AggregateFunction,
) -> Result<SqlValue, ExecutorError> {
    // Get column from batch
    let schema = batch.schema();
    let (col_idx, _) = schema.column_with_name(column_name).ok_or_else(|| {
        ExecutorError::ColumnarColumnNotFound {
            column_index: 0, // Column referenced by name
            batch_columns: schema.fields().len(),
        }
    })?;

    let column = batch.column(col_idx);

    // Dispatch to type-specific SIMD aggregation
    match function {
        AggregateFunction::Count => {
            // COUNT is trivial: number of non-null values
            let non_null_count = column.len() - column.null_count();
            Ok(SqlValue::Integer(non_null_count as i64))
        }
        AggregateFunction::Sum => aggregate_sum_simd(column),
        AggregateFunction::Avg => aggregate_avg_simd(column),
        AggregateFunction::Min => aggregate_min_simd(column),
        AggregateFunction::Max => aggregate_max_simd(column),
    }
}

/// SIMD-accelerated SUM aggregation
fn aggregate_sum_simd(column: &ArrayRef) -> Result<SqlValue, ExecutorError> {
    match column.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let array = column.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "aggregate_sum_simd".to_string(),
                }
            })?;

            let result = sum(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "SUM".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Integer(result))
        }
        arrow::datatypes::DataType::Float64 => {
            let array = column.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "aggregate_sum_simd".to_string(),
                }
            })?;

            let result = sum(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "SUM".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Double(result))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "SUM".to_string(),
            array_type: format!("{:?}", column.data_type()),
        }),
    }
}

/// SIMD-accelerated AVG aggregation
fn aggregate_avg_simd(column: &ArrayRef) -> Result<SqlValue, ExecutorError> {
    match column.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let array = column.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "aggregate_avg_simd".to_string(),
                }
            })?;

            let sum_val = sum(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "AVG (SUM)".to_string(),
                reason: "returned None".to_string(),
            })?;

            let count = (array.len() - array.null_count()) as f64;
            if count == 0.0 {
                return Ok(SqlValue::Null);
            }

            Ok(SqlValue::Double(sum_val as f64 / count))
        }
        arrow::datatypes::DataType::Float64 => {
            let array = column.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "aggregate_avg_simd".to_string(),
                }
            })?;

            let sum_val = sum(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "AVG (SUM)".to_string(),
                reason: "returned None".to_string(),
            })?;

            let count = (array.len() - array.null_count()) as f64;
            if count == 0.0 {
                return Ok(SqlValue::Null);
            }

            Ok(SqlValue::Double(sum_val / count))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "AVG".to_string(),
            array_type: format!("{:?}", column.data_type()),
        }),
    }
}

/// SIMD-accelerated MIN aggregation
fn aggregate_min_simd(column: &ArrayRef) -> Result<SqlValue, ExecutorError> {
    match column.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let array = column.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "aggregate_min_simd".to_string(),
                }
            })?;

            let result = min(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MIN".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Integer(result))
        }
        arrow::datatypes::DataType::Float64 => {
            let array = column.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "aggregate_min_simd".to_string(),
                }
            })?;

            let result = min(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MIN".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Double(result))
        }
        arrow::datatypes::DataType::Date32 => {
            use super::batch::days_since_epoch_to_date;

            let array = column.as_any().downcast_ref::<Date32Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Date32Array".to_string(),
                    context: "aggregate_min_simd".to_string(),
                }
            })?;

            let result = min(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MIN".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Date(days_since_epoch_to_date(result)))
        }
        arrow::datatypes::DataType::Timestamp(TimeUnit::Microsecond, None) => {
            use super::batch::microseconds_to_timestamp;

            let array =
                column.as_any().downcast_ref::<TimestampMicrosecondArray>().ok_or_else(|| {
                    ExecutorError::ArrowDowncastError {
                        expected_type: "TimestampMicrosecondArray".to_string(),
                        context: "aggregate_min_simd".to_string(),
                    }
                })?;

            let result = min(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MIN".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Timestamp(microseconds_to_timestamp(result)))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "MIN".to_string(),
            array_type: format!("{:?}", column.data_type()),
        }),
    }
}

/// SIMD-accelerated MAX aggregation
fn aggregate_max_simd(column: &ArrayRef) -> Result<SqlValue, ExecutorError> {
    match column.data_type() {
        arrow::datatypes::DataType::Int64 => {
            let array = column.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Int64Array".to_string(),
                    context: "aggregate_max_simd".to_string(),
                }
            })?;

            let result = max(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MAX".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Integer(result))
        }
        arrow::datatypes::DataType::Float64 => {
            let array = column.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Float64Array".to_string(),
                    context: "aggregate_max_simd".to_string(),
                }
            })?;

            let result = max(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MAX".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Double(result))
        }
        arrow::datatypes::DataType::Date32 => {
            use super::batch::days_since_epoch_to_date;

            let array = column.as_any().downcast_ref::<Date32Array>().ok_or_else(|| {
                ExecutorError::ArrowDowncastError {
                    expected_type: "Date32Array".to_string(),
                    context: "aggregate_max_simd".to_string(),
                }
            })?;

            let result = max(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MAX".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Date(days_since_epoch_to_date(result)))
        }
        arrow::datatypes::DataType::Timestamp(TimeUnit::Microsecond, None) => {
            use super::batch::microseconds_to_timestamp;

            let array =
                column.as_any().downcast_ref::<TimestampMicrosecondArray>().ok_or_else(|| {
                    ExecutorError::ArrowDowncastError {
                        expected_type: "TimestampMicrosecondArray".to_string(),
                        context: "aggregate_max_simd".to_string(),
                    }
                })?;

            let result = max(array).ok_or_else(|| ExecutorError::SimdOperationFailed {
                operation: "MAX".to_string(),
                reason: "returned None (all nulls?)".to_string(),
            })?;

            Ok(SqlValue::Timestamp(microseconds_to_timestamp(result)))
        }
        _ => Err(ExecutorError::UnsupportedArrayType {
            operation: "MAX".to_string(),
            array_type: format!("{:?}", column.data_type()),
        }),
    }
}

/// Perform multiple aggregations on a RecordBatch
pub fn aggregate_batch_simd(
    batch: &RecordBatch,
    aggregates: &[(String, AggregateFunction)],
) -> Result<Vec<SqlValue>, ExecutorError> {
    let mut results = Vec::with_capacity(aggregates.len());

    for (column_name, function) in aggregates {
        let value = aggregate_column_simd(batch, column_name, function.clone())?;
        results.push(value);
    }

    Ok(results)
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow::datatypes::{DataType, Field, Schema};

    use super::*;

    #[test]
    fn test_simd_sum_int64() {
        // Create test data: [1, 2, 3, 4, 5]
        let schema = Schema::new(vec![Field::new("value", DataType::Int64, false)]);
        let array = Int64Array::from(vec![1, 2, 3, 4, 5]);
        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();

        let result = aggregate_column_simd(&batch, "value", AggregateFunction::Sum).unwrap();

        assert_eq!(result, SqlValue::Integer(15)); // 1+2+3+4+5 = 15
    }

    #[test]
    fn test_simd_avg_float64() {
        // Create test data: [1.0, 2.0, 3.0, 4.0, 5.0]
        let schema = Schema::new(vec![Field::new("value", DataType::Float64, false)]);
        let array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();

        let result = aggregate_column_simd(&batch, "value", AggregateFunction::Avg).unwrap();

        match result {
            SqlValue::Double(f) => assert!((f - 3.0).abs() < 0.0001), // avg = 3.0
            _ => panic!("Expected Double value"),
        }
    }

    #[test]
    fn test_simd_min_max() {
        // Create test data: [5, 2, 8, 1, 9]
        let schema = Schema::new(vec![Field::new("value", DataType::Int64, false)]);
        let array = Int64Array::from(vec![5, 2, 8, 1, 9]);
        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();

        let min_result = aggregate_column_simd(&batch, "value", AggregateFunction::Min).unwrap();
        let max_result = aggregate_column_simd(&batch, "value", AggregateFunction::Max).unwrap();

        assert_eq!(min_result, SqlValue::Integer(1));
        assert_eq!(max_result, SqlValue::Integer(9));
    }
}