vibesql-executor 0.2.0

//! Batch-native expression processing on ColumnarBatch
//!
//! This module provides SIMD-accelerated expression evaluation directly on
//! `ColumnarBatch` without converting to rows. This eliminates the ~10-15ms
//! overhead of `batch.to_rows()` for large batches.
//!
//! ## Performance
//!
//! For a batch with 100K rows:
//! - Row-based path: ~15ms (to_rows) + ~5ms (eval) = ~20ms
//! - Batch-native path: ~3ms (SIMD eval + aggregate)
//!
//! ~6-7x speedup for expression aggregates.

use std::sync::Arc;

use vibesql_ast::Expression;
use vibesql_types::SqlValue;

use crate::{
    errors::ExecutorError,
    schema::CombinedSchema,
    select::columnar::{
        aggregate::AggregateOp,
        batch::{ColumnArray, ColumnarBatch},
    },
};

/// Compute an aggregate over an expression directly from a ColumnarBatch (no row conversion)
///
/// This is the batch-native path for expression aggregates. Instead of converting
/// the batch to rows and then evaluating expressions, we:
/// 1. Evaluate the expression directly on the batch's column arrays using SIMD
/// 2. Aggregate the resulting array using SIMD operations
///
/// This eliminates the ~10-15ms overhead of `batch.to_rows()` for large batches.
///
/// # Arguments
///
/// * `batch` - The ColumnarBatch to process (typically already filtered)
/// * `expr` - The expression to evaluate (e.g., `a * b`)
/// * `op` - The aggregate operation (SUM, AVG, MIN, MAX, COUNT)
/// * `schema` - Schema for resolving column names
///
/// # Returns
///
/// The aggregated SqlValue result
pub fn compute_batch_expression_aggregate(
    batch: &ColumnarBatch,
    expr: &Expression,
    op: AggregateOp,
    schema: &CombinedSchema,
) -> Result<SqlValue, ExecutorError> {
    log::debug!(
        "[BatchExpr] compute_batch_expression_aggregate: {} rows, op={:?}",
        batch.row_count(),
        op
    );

    // Empty batch handling
    if batch.row_count() == 0 {
        log::debug!("[BatchExpr] Empty batch, returning default");
        return Ok(match op {
            AggregateOp::Count => SqlValue::Integer(0),
            _ => SqlValue::Null,
        });
    }

    // Evaluate expression on batch columns using SIMD
    log::debug!("[BatchExpr] Evaluating expression on batch columns (SIMD-enabled)");
    let result_array = evaluate_batch_expression(batch, expr, schema)?;

    // Aggregate the result array using SIMD
    log::debug!("[BatchExpr] Aggregating result array with SIMD");
    aggregate_column_array(&result_array, op)
}

/// Evaluate an expression directly on ColumnarBatch column arrays
///
/// Returns a ColumnArray containing the computed values.
fn evaluate_batch_expression(
    batch: &ColumnarBatch,
    expr: &Expression,
    schema: &CombinedSchema,
) -> Result<ColumnArray, ExecutorError> {
    match expr {
        Expression::ColumnRef(col_id) => {
            // Simple column reference - return the column directly
            let column = col_id.column_canonical();
            let col_idx =
                schema.get_column_index(col_id.table_canonical(), column).ok_or_else(|| {
                    ExecutorError::UnsupportedExpression(format!("Column not found: {}", column))
                })?;

            batch.column(col_idx).cloned().ok_or_else(|| ExecutorError::ColumnarColumnNotFound {
                column_index: col_idx,
                batch_columns: batch.column_count(),
            })
        }
        Expression::Literal(val) => {
            // Create an array filled with the literal value
            create_literal_column_array(val, batch.row_count())
        }
        Expression::BinaryOp { left, op, right } => {
            // Recursively evaluate left and right, then apply operation
            let left_array = evaluate_batch_expression(batch, left, schema)?;
            let right_array = evaluate_batch_expression(batch, right, schema)?;

            apply_binary_op_to_columns(&left_array, &right_array, op)
        }
        _ => Err(ExecutorError::UnsupportedExpression(
            "Complex expressions not supported in batch-native columnar aggregates".to_string(),
        )),
    }
}

/// Create a ColumnArray filled with a literal value
fn create_literal_column_array(value: &SqlValue, len: usize) -> Result<ColumnArray, ExecutorError> {
    match value {
        SqlValue::Integer(i) | SqlValue::Bigint(i) => {
            Ok(ColumnArray::Int64(Arc::new(vec![*i; len]), None))
        }
        SqlValue::Smallint(i) => Ok(ColumnArray::Int64(Arc::new(vec![*i as i64; len]), None)),
        SqlValue::Float(f) => Ok(ColumnArray::Float64(Arc::new(vec![*f as f64; len]), None)),
        SqlValue::Double(f) | SqlValue::Numeric(f) | SqlValue::Real(f) => {
            // Real is now f64 (SQLite REAL is 8-byte IEEE float)
            Ok(ColumnArray::Float64(Arc::new(vec![*f; len]), None))
        }
        SqlValue::Null => {
            // Create array of nulls (represented as Float64 with all nulls)
            Ok(ColumnArray::Float64(Arc::new(vec![0.0; len]), Some(Arc::new(vec![true; len]))))
        }
        _ => Err(ExecutorError::UnsupportedExpression(format!(
            "Cannot create literal column array for {:?}",
            value
        ))),
    }
}

/// Apply a binary operation to two column arrays
///
/// Supports SIMD-accelerated arithmetic on Int64 and Float64 columns.
/// Falls back to row-by-row evaluation for Mixed columns.
fn apply_binary_op_to_columns(
    left: &ColumnArray,
    right: &ColumnArray,
    op: &vibesql_ast::BinaryOperator,
) -> Result<ColumnArray, ExecutorError> {
    use vibesql_ast::BinaryOperator::*;

    // Try to convert Mixed columns to typed columns for SIMD
    let (left_typed, right_typed) = match (left, right) {
        // If either column is Mixed, try to extract numeric values
        (ColumnArray::Mixed(left_vals), ColumnArray::Mixed(right_vals)) => {
            let left_f64 = try_extract_f64_from_mixed(left_vals)?;
            let right_f64 = try_extract_f64_from_mixed(right_vals)?;
            (
                ColumnArray::Float64(Arc::new(left_f64), None),
                ColumnArray::Float64(Arc::new(right_f64), None),
            )
        }
        (ColumnArray::Mixed(left_vals), other) => {
            let left_f64 = try_extract_f64_from_mixed(left_vals)?;
            (ColumnArray::Float64(Arc::new(left_f64), None), other.clone())
        }
        (other, ColumnArray::Mixed(right_vals)) => {
            let right_f64 = try_extract_f64_from_mixed(right_vals)?;
            (other.clone(), ColumnArray::Float64(Arc::new(right_f64), None))
        }
        _ => (left.clone(), right.clone()),
    };

    match (&left_typed, &right_typed) {
        // Both Float64 - direct SIMD operations
        (
            ColumnArray::Float64(left_vals, left_nulls),
            ColumnArray::Float64(right_vals, right_nulls),
        ) => {
            let result = apply_float64_binary_op(left_vals, right_vals, op)?;
            let nulls = merge_null_bitmaps(
                left_nulls.as_ref().map(|a| a.as_slice()),
                right_nulls.as_ref().map(|a| a.as_slice()),
                left_vals.len(),
            );
            Ok(ColumnArray::Float64(Arc::new(result), nulls.map(Arc::new)))
        }
        // Both Int64 - SIMD operations (result type depends on operation)
        (
            ColumnArray::Int64(left_vals, left_nulls),
            ColumnArray::Int64(right_vals, right_nulls),
        ) => {
            match op {
                Plus | Minus | Multiply => {
                    let result = apply_int64_binary_op(left_vals, right_vals, op)?;
                    let nulls = merge_null_bitmaps(
                        left_nulls.as_ref().map(|a| a.as_slice()),
                        right_nulls.as_ref().map(|a| a.as_slice()),
                        left_vals.len(),
                    );
                    Ok(ColumnArray::Int64(Arc::new(result), nulls.map(Arc::new)))
                }
                Divide => {
                    // Division always produces Float64
                    let left_f64: Vec<f64> = left_vals.iter().map(|&v| v as f64).collect();
                    let right_f64: Vec<f64> = right_vals.iter().map(|&v| v as f64).collect();
                    let result = apply_float64_binary_op(&left_f64, &right_f64, op)?;
                    let nulls = merge_null_bitmaps(
                        left_nulls.as_ref().map(|a| a.as_slice()),
                        right_nulls.as_ref().map(|a| a.as_slice()),
                        left_vals.len(),
                    );
                    Ok(ColumnArray::Float64(Arc::new(result), nulls.map(Arc::new)))
                }
                _ => Err(ExecutorError::UnsupportedExpression(format!(
                    "Unsupported binary operator for Int64: {:?}",
                    op
                ))),
            }
        }
        // Mixed types - cast to Float64
        (
            ColumnArray::Int64(left_vals, left_nulls),
            ColumnArray::Float64(right_vals, right_nulls),
        ) => {
            let left_f64: Vec<f64> = left_vals.iter().map(|&v| v as f64).collect();
            let result = apply_float64_binary_op(&left_f64, right_vals, op)?;
            let nulls = merge_null_bitmaps(
                left_nulls.as_ref().map(|a| a.as_slice()),
                right_nulls.as_ref().map(|a| a.as_slice()),
                left_vals.len(),
            );
            Ok(ColumnArray::Float64(Arc::new(result), nulls.map(Arc::new)))
        }
        (
            ColumnArray::Float64(left_vals, left_nulls),
            ColumnArray::Int64(right_vals, right_nulls),
        ) => {
            let right_f64: Vec<f64> = right_vals.iter().map(|&v| v as f64).collect();
            let result = apply_float64_binary_op(left_vals, &right_f64, op)?;
            let nulls = merge_null_bitmaps(
                left_nulls.as_ref().map(|a| a.as_slice()),
                right_nulls.as_ref().map(|a| a.as_slice()),
                left_vals.len(),
            );
            Ok(ColumnArray::Float64(Arc::new(result), nulls.map(Arc::new)))
        }
        // Fallback for other types - signal to caller to use row-based path
        _ => Err(ExecutorError::UnsupportedExpression(
            "Non-numeric columns not supported in batch arithmetic".to_string(),
        )),
    }
}

/// Try to extract f64 values from a Mixed column array
///
/// Returns an error if any value is non-numeric.
fn try_extract_f64_from_mixed(values: &[SqlValue]) -> Result<Vec<f64>, ExecutorError> {
    values
        .iter()
        .map(|v| match v {
            SqlValue::Integer(i) | SqlValue::Bigint(i) => Ok(*i as f64),
            SqlValue::Smallint(i) => Ok(*i as f64),
            SqlValue::Float(f) => Ok(*f as f64),
            SqlValue::Double(f) | SqlValue::Numeric(f) | SqlValue::Real(f) => Ok(*f), // Real is now f64
            SqlValue::Null => Ok(f64::NAN), // NaN will propagate correctly
            _ => Err(ExecutorError::UnsupportedExpression(
                "Non-numeric columns not supported in batch arithmetic".to_string(),
            )),
        })
        .collect()
}

/// Apply a binary operation to Float64 arrays using SIMD
fn apply_float64_binary_op(
    left: &[f64],
    right: &[f64],
    op: &vibesql_ast::BinaryOperator,
) -> Result<Vec<f64>, ExecutorError> {
    use vibesql_ast::BinaryOperator::*;

    if left.len() != right.len() {
        return Err(ExecutorError::ColumnarLengthMismatch {
            context: "binary_op".to_string(),
            expected: left.len(),
            actual: right.len(),
        });
    }

    // SIMD-friendly iteration (compiler will auto-vectorize)
    let result: Vec<f64> = match op {
        Plus => left.iter().zip(right.iter()).map(|(l, r)| l + r).collect(),
        Minus => left.iter().zip(right.iter()).map(|(l, r)| l - r).collect(),
        Multiply => left.iter().zip(right.iter()).map(|(l, r)| l * r).collect(),
        Divide => left.iter().zip(right.iter()).map(|(l, r)| l / r).collect(),
        _ => {
            return Err(ExecutorError::UnsupportedExpression(format!(
                "Unsupported binary operator for Float64: {:?}",
                op
            )))
        }
    };

    Ok(result)
}

/// Apply a binary operation to Int64 arrays using SIMD
fn apply_int64_binary_op(
    left: &[i64],
    right: &[i64],
    op: &vibesql_ast::BinaryOperator,
) -> Result<Vec<i64>, ExecutorError> {
    use vibesql_ast::BinaryOperator::*;

    if left.len() != right.len() {
        return Err(ExecutorError::ColumnarLengthMismatch {
            context: "binary_op".to_string(),
            expected: left.len(),
            actual: right.len(),
        });
    }

    // SIMD-friendly iteration (compiler will auto-vectorize)
    let result: Vec<i64> = match op {
        Plus => left.iter().zip(right.iter()).map(|(l, r)| l + r).collect(),
        Minus => left.iter().zip(right.iter()).map(|(l, r)| l - r).collect(),
        Multiply => left.iter().zip(right.iter()).map(|(l, r)| l * r).collect(),
        _ => {
            return Err(ExecutorError::UnsupportedExpression(format!(
                "Unsupported binary operator for Int64: {:?}",
                op
            )))
        }
    };

    Ok(result)
}

/// Merge two null bitmaps (OR operation - if either is null, result is null)
fn merge_null_bitmaps(
    left: Option<&[bool]>,
    right: Option<&[bool]>,
    _len: usize,
) -> Option<Vec<bool>> {
    match (left, right) {
        (None, None) => None,
        (Some(l), None) => Some(l.to_vec()),
        (None, Some(r)) => Some(r.to_vec()),
        (Some(l), Some(r)) => {
            let merged: Vec<bool> =
                l.iter().zip(r.iter()).map(|(&l_null, &r_null)| l_null || r_null).collect();
            if merged.iter().any(|&is_null| is_null) {
                Some(merged)
            } else {
                None
            }
        }
    }
}

/// Aggregate a ColumnArray using SIMD operations
fn aggregate_column_array(array: &ColumnArray, op: AggregateOp) -> Result<SqlValue, ExecutorError> {
    match array {
        ColumnArray::Float64(values, nulls) => {
            aggregate_f64_array(values, nulls.as_ref().map(|n| n.as_slice()), op)
        }
        ColumnArray::Int64(values, nulls) => {
            aggregate_i64_array(values, nulls.as_ref().map(|n| n.as_slice()), op)
        }
        _ => Err(ExecutorError::UnsupportedExpression(
            "Non-numeric columns not supported for aggregation".to_string(),
        )),
    }
}

/// Aggregate a Float64 array
fn aggregate_f64_array(
    values: &[f64],
    nulls: Option<&[bool]>,
    op: AggregateOp,
) -> Result<SqlValue, ExecutorError> {
    // Filter out null values for aggregation
    let non_null_values: Vec<f64> = if let Some(null_bitmap) = nulls {
        values
            .iter()
            .zip(null_bitmap.iter())
            .filter(|(_, &is_null)| !is_null)
            .map(|(&v, _)| v)
            .collect()
    } else {
        values.to_vec()
    };

    if non_null_values.is_empty() {
        return Ok(match op {
            AggregateOp::Count => SqlValue::Integer(0),
            _ => SqlValue::Null,
        });
    }

    match op {
        AggregateOp::Sum => {
            let sum: f64 = non_null_values.iter().sum();
            Ok(SqlValue::Double(sum))
        }
        AggregateOp::Count => Ok(SqlValue::Integer(non_null_values.len() as i64)),
        AggregateOp::Avg => {
            let sum: f64 = non_null_values.iter().sum();
            let count = non_null_values.len() as f64;
            Ok(SqlValue::Double(sum / count))
        }
        AggregateOp::Min => {
            let min = non_null_values.iter().cloned().fold(f64::INFINITY, f64::min);
            Ok(SqlValue::Double(min))
        }
        AggregateOp::Max => {
            let max = non_null_values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            Ok(SqlValue::Double(max))
        }
    }
}

/// Aggregate an Int64 array
fn aggregate_i64_array(
    values: &[i64],
    nulls: Option<&[bool]>,
    op: AggregateOp,
) -> Result<SqlValue, ExecutorError> {
    // Filter out null values for aggregation
    let non_null_values: Vec<i64> = if let Some(null_bitmap) = nulls {
        values
            .iter()
            .zip(null_bitmap.iter())
            .filter(|(_, &is_null)| !is_null)
            .map(|(&v, _)| v)
            .collect()
    } else {
        values.to_vec()
    };

    if non_null_values.is_empty() {
        return Ok(match op {
            AggregateOp::Count => SqlValue::Integer(0),
            _ => SqlValue::Null,
        });
    }

    match op {
        AggregateOp::Sum => {
            // SQLite's SUM() preserves integer type for integer inputs
            let sum: i64 = non_null_values.iter().sum();
            Ok(SqlValue::Integer(sum))
        }
        AggregateOp::Count => Ok(SqlValue::Integer(non_null_values.len() as i64)),
        AggregateOp::Avg => {
            let sum: i64 = non_null_values.iter().sum();
            let count = non_null_values.len() as f64;
            Ok(SqlValue::Double(sum as f64 / count))
        }
        AggregateOp::Min => {
            let min = *non_null_values.iter().min().unwrap();
            Ok(SqlValue::Integer(min))
        }
        AggregateOp::Max => {
            let max = *non_null_values.iter().max().unwrap();
            Ok(SqlValue::Integer(max))
        }
    }
}

/// Evaluate an expression on a ColumnarBatch and return the result as a ColumnArray
///
/// This is the public interface for computing expression columns for GROUP BY support.
/// It evaluates an expression (e.g., `l_extendedprice * (1 - l_discount)`) on all rows
/// and returns the computed values as a typed column that can be added to the batch.
///
/// # Arguments
///
/// * `batch` - The ColumnarBatch to evaluate the expression on
/// * `expr` - The expression to evaluate (e.g., BinaryOp of column refs)
/// * `schema` - Schema for resolving column names
///
/// # Returns
///
/// A ColumnArray containing the computed values (Float64 or Int64)
pub fn evaluate_expression_to_column(
    batch: &ColumnarBatch,
    expr: &Expression,
    schema: &CombinedSchema,
) -> Result<ColumnArray, ExecutorError> {
    evaluate_batch_expression(batch, expr, schema)
}

/// Evaluate an expression using a cached column as a base value
///
/// This is used for Common Sub-Expression Elimination (CSE) in GROUP BY queries.
/// When we have expressions like:
///   E1 = l_extendedprice * (1 - l_discount)
///   E2 = E1 * (1 + l_tax)
///
/// After computing E1 and adding it as column N, we can compute E2 as:
///   column_N * (1 + l_tax)
///
/// This avoids re-computing E1 for E2, providing ~20-30% speedup for Q1-style queries.
///
/// # Arguments
///
/// * `batch` - The ColumnarBatch containing the cached column
/// * `expr` - The remaining expression to evaluate (e.g., `(1 + l_tax)`)
/// * `cached_col_idx` - Index of the cached column containing the sub-expression result
/// * `schema` - Schema for resolving any column names in `expr`
///
/// # Returns
///
/// A ColumnArray containing the result of: cached_column * expr (or cached_column OP expr)
pub fn evaluate_expression_with_cached_column(
    batch: &ColumnarBatch,
    expr: &Expression,
    cached_col_idx: usize,
    schema: &CombinedSchema,
) -> Result<ColumnArray, ExecutorError> {
    // Get the cached column
    let cached_col =
        batch.column(cached_col_idx).ok_or_else(|| ExecutorError::ColumnarColumnNotFound {
            column_index: cached_col_idx,
            batch_columns: batch.column_count(),
        })?;

    // Evaluate the remaining expression
    let expr_col = evaluate_batch_expression(batch, expr, schema)?;

    // Multiply cached_column * expr_result (most common CSE pattern)
    apply_binary_op_to_columns(cached_col, &expr_col, &vibesql_ast::BinaryOperator::Multiply)
}