featherdb_query/

window.rs

//! Window function execution
//!
//! Implements SQL window functions (OVER clause) including:
//! - Ranking functions: ROW_NUMBER, RANK, DENSE_RANK, NTILE
//! - Navigation functions: LAG, LEAD, FIRST_VALUE, LAST_VALUE, NTH_VALUE
//! - Aggregate functions over windows: SUM, AVG, COUNT, MIN, MAX

use crate::executor::Row;
use crate::expr::{
    Expr, FrameBound, WindowFrame, WindowFunction, WindowFunctionType, WindowOrderByExpr,
    WindowSortOrder,
};
use featherdb_core::{Result, Value};
use std::collections::HashMap;
use std::sync::Arc;

/// Window executor that computes window function results
pub struct WindowExecutor;

impl WindowExecutor {
    /// Execute window functions on a set of rows
    ///
    /// For each window expression, computes the window function result for each row
    /// and appends it to the row's values.
    pub fn execute(rows: Vec<Row>, window_exprs: &[(Expr, String)]) -> Result<Vec<Row>> {
        if rows.is_empty() {
            return Ok(rows);
        }

        // Build column map from first row
        let col_map = build_col_map(&rows[0].columns);

        // Process each window expression
        let mut result_rows = rows;

        for (window_expr, alias) in window_exprs {
            if let Expr::Window(window_func) = window_expr {
                result_rows =
                    Self::execute_window_function(result_rows, window_func, alias, &col_map)?;
            }
        }

        Ok(result_rows)
    }

    /// Execute a single window function across all rows
    fn execute_window_function(
        mut rows: Vec<Row>,
        window_func: &WindowFunction,
        alias: &str,
        col_map: &HashMap<String, usize>,
    ) -> Result<Vec<Row>> {
        // Partition the rows
        let partitions = Self::partition_rows(&rows, &window_func.partition_by, col_map)?;

        // For each partition, sort and compute window values
        let mut window_values: Vec<(usize, Value)> = Vec::new();

        for partition_indices in partitions {
            // Get a view of partition rows with their original indices
            let mut partition_rows: Vec<(usize, &Row)> =
                partition_indices.iter().map(|&i| (i, &rows[i])).collect();

            // Sort the partition by ORDER BY clause
            if !window_func.order_by.is_empty() {
                Self::sort_partition(&mut partition_rows, &window_func.order_by, col_map)?;
            }

            // Compute window function values for this partition
            let values = Self::compute_window_values(
                &partition_rows,
                &window_func.function,
                window_func.frame.as_ref(),
                &window_func.order_by,
                col_map,
            )?;

            // Map values back to original row indices
            for ((orig_idx, _), value) in partition_rows.into_iter().zip(values) {
                window_values.push((orig_idx, value));
            }
        }

        // Sort window values by original index
        window_values.sort_by_key(|(idx, _)| *idx);

        // Add window values to rows
        for (row, (_, value)) in rows.iter_mut().zip(window_values) {
            row.values.push(value);
            // Need to make columns mutable via Arc::make_mut
            let cols = Arc::make_mut(&mut row.columns);
            cols.push(alias.to_string());
        }

        Ok(rows)
    }

    /// Partition rows based on PARTITION BY expressions
    fn partition_rows(
        rows: &[Row],
        partition_by: &[Expr],
        col_map: &HashMap<String, usize>,
    ) -> Result<Vec<Vec<usize>>> {
        if partition_by.is_empty() {
            // No partitioning - all rows in one partition
            return Ok(vec![(0..rows.len()).collect()]);
        }

        let mut partitions: HashMap<Vec<Value>, Vec<usize>> = HashMap::new();

        for (idx, row) in rows.iter().enumerate() {
            let key: Vec<Value> = partition_by
                .iter()
                .map(|expr| expr.eval(&row.values, col_map))
                .collect::<Result<_>>()?;

            partitions.entry(key).or_default().push(idx);
        }

        Ok(partitions.into_values().collect())
    }

    /// Sort partition rows by ORDER BY expressions
    fn sort_partition(
        partition: &mut [(usize, &Row)],
        order_by: &[WindowOrderByExpr],
        col_map: &HashMap<String, usize>,
    ) -> Result<()> {
        partition.sort_by(|(_, row_a), (_, row_b)| {
            for order_expr in order_by {
                let val_a = order_expr
                    .expr
                    .eval(&row_a.values, col_map)
                    .unwrap_or(Value::Null);
                let val_b = order_expr
                    .expr
                    .eval(&row_b.values, col_map)
                    .unwrap_or(Value::Null);

                let cmp = match order_expr.order {
                    WindowSortOrder::Asc => val_a.cmp(&val_b),
                    WindowSortOrder::Desc => val_b.cmp(&val_a),
                };

                if cmp != std::cmp::Ordering::Equal {
                    return cmp;
                }
            }
            std::cmp::Ordering::Equal
        });

        Ok(())
    }

    /// Compute window function values for a sorted partition
    fn compute_window_values(
        partition: &[(usize, &Row)],
        function: &WindowFunctionType,
        frame: Option<&WindowFrame>,
        order_by: &[WindowOrderByExpr],
        col_map: &HashMap<String, usize>,
    ) -> Result<Vec<Value>> {
        let n = partition.len();
        let mut values = Vec::with_capacity(n);

        for (pos, (_, row)) in partition.iter().enumerate() {
            let value = match function {
                WindowFunctionType::RowNumber => Value::Integer((pos + 1) as i64),

                WindowFunctionType::Rank => {
                    Self::compute_rank(partition, pos, order_by, col_map, false)?
                }

                WindowFunctionType::DenseRank => {
                    Self::compute_rank(partition, pos, order_by, col_map, true)?
                }

                WindowFunctionType::NTile(num_buckets) => {
                    let bucket = ((pos as u32 * *num_buckets) / n as u32) + 1;
                    Value::Integer(bucket as i64)
                }

                WindowFunctionType::Lag {
                    expr,
                    offset,
                    default,
                } => {
                    let target_pos = pos as i64 - *offset;
                    if target_pos >= 0 && (target_pos as usize) < n {
                        let target_row = partition[target_pos as usize].1;
                        expr.eval(&target_row.values, col_map)?
                    } else {
                        default
                            .as_ref()
                            .map(|d| d.eval(&row.values, col_map))
                            .transpose()?
                            .unwrap_or(Value::Null)
                    }
                }

                WindowFunctionType::Lead {
                    expr,
                    offset,
                    default,
                } => {
                    let target_pos = pos as i64 + *offset;
                    if target_pos >= 0 && (target_pos as usize) < n {
                        let target_row = partition[target_pos as usize].1;
                        expr.eval(&target_row.values, col_map)?
                    } else {
                        default
                            .as_ref()
                            .map(|d| d.eval(&row.values, col_map))
                            .transpose()?
                            .unwrap_or(Value::Null)
                    }
                }

                WindowFunctionType::FirstValue(expr) => {
                    let (start, _) = Self::compute_frame_bounds(pos, n, frame);
                    if start < n {
                        let first_row = partition[start].1;
                        expr.eval(&first_row.values, col_map)?
                    } else {
                        Value::Null
                    }
                }

                WindowFunctionType::LastValue(expr) => {
                    let (_, end) = Self::compute_frame_bounds(pos, n, frame);
                    if end > 0 && end <= n {
                        let last_row = partition[end - 1].1;
                        expr.eval(&last_row.values, col_map)?
                    } else {
                        Value::Null
                    }
                }

                WindowFunctionType::NthValue(expr, nth) => {
                    let (start, end) = Self::compute_frame_bounds(pos, n, frame);
                    let target = start + (*nth as usize) - 1;
                    if target < end && target < n {
                        let target_row = partition[target].1;
                        expr.eval(&target_row.values, col_map)?
                    } else {
                        Value::Null
                    }
                }

                // Aggregate functions over window
                WindowFunctionType::Sum(expr) => {
                    Self::compute_aggregate_sum(partition, pos, expr, frame, col_map)?
                }

                WindowFunctionType::Avg(expr) => {
                    Self::compute_aggregate_avg(partition, pos, expr, frame, col_map)?
                }

                WindowFunctionType::Count(expr) => {
                    Self::compute_aggregate_count(partition, pos, expr.as_deref(), frame, col_map)?
                }

                WindowFunctionType::Min(expr) => {
                    Self::compute_aggregate_min(partition, pos, expr, frame, col_map)?
                }

                WindowFunctionType::Max(expr) => {
                    Self::compute_aggregate_max(partition, pos, expr, frame, col_map)?
                }
            };

            values.push(value);
        }

        Ok(values)
    }

    /// Compare ORDER BY values of two rows for tie detection
    fn order_by_values_equal(
        row_a: &Row,
        row_b: &Row,
        order_by: &[WindowOrderByExpr],
        col_map: &HashMap<String, usize>,
    ) -> bool {
        for order_expr in order_by {
            let val_a = order_expr
                .expr
                .eval(&row_a.values, col_map)
                .unwrap_or(Value::Null);
            let val_b = order_expr
                .expr
                .eval(&row_b.values, col_map)
                .unwrap_or(Value::Null);
            if val_a != val_b {
                return false;
            }
        }
        true
    }

    /// Compute RANK or DENSE_RANK with proper tie detection
    fn compute_rank(
        partition: &[(usize, &Row)],
        pos: usize,
        order_by: &[WindowOrderByExpr],
        col_map: &HashMap<String, usize>,
        dense: bool,
    ) -> Result<Value> {
        if pos == 0 || order_by.is_empty() {
            return Ok(Value::Integer(1));
        }

        // Walk backwards to find the rank
        if dense {
            // DENSE_RANK: count distinct ORDER BY value groups before this position
            let mut dense_rank = 1i64;
            for i in 1..=pos {
                if !Self::order_by_values_equal(
                    partition[i].1,
                    partition[i - 1].1,
                    order_by,
                    col_map,
                ) {
                    dense_rank += 1;
                }
            }
            Ok(Value::Integer(dense_rank))
        } else {
            // RANK: position of first row with same ORDER BY values (1-based)
            let mut rank_start = pos;
            while rank_start > 0
                && Self::order_by_values_equal(
                    partition[rank_start].1,
                    partition[rank_start - 1].1,
                    order_by,
                    col_map,
                )
            {
                rank_start -= 1;
            }
            Ok(Value::Integer((rank_start + 1) as i64))
        }
    }

    /// Compute frame bounds for the current position
    fn compute_frame_bounds(
        pos: usize,
        partition_size: usize,
        frame: Option<&WindowFrame>,
    ) -> (usize, usize) {
        let frame = match frame {
            Some(f) => f,
            None => {
                // Default frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
                return (0, pos + 1);
            }
        };

        let start = match &frame.start {
            FrameBound::UnboundedPreceding => 0,
            FrameBound::Preceding(n) => pos.saturating_sub(*n as usize),
            FrameBound::CurrentRow => pos,
            FrameBound::Following(n) => (pos + *n as usize).min(partition_size),
            FrameBound::UnboundedFollowing => partition_size,
        };

        let end = match &frame.end {
            FrameBound::UnboundedPreceding => 0,
            FrameBound::Preceding(n) => pos.saturating_sub(*n as usize),
            FrameBound::CurrentRow => pos + 1,
            FrameBound::Following(n) => (pos + *n as usize + 1).min(partition_size),
            FrameBound::UnboundedFollowing => partition_size,
        };

        (start, end)
    }

    /// Compute SUM aggregate over window frame
    fn compute_aggregate_sum(
        partition: &[(usize, &Row)],
        pos: usize,
        expr: &Expr,
        frame: Option<&WindowFrame>,
        col_map: &HashMap<String, usize>,
    ) -> Result<Value> {
        let (start, end) = Self::compute_frame_bounds(pos, partition.len(), frame);
        let mut sum = 0.0;
        let mut has_value = false;

        for (_, row) in partition.iter().skip(start).take(end - start) {
            let val = expr.eval(&row.values, col_map)?;
            if let Some(n) = val.as_f64() {
                sum += n;
                has_value = true;
            }
        }

        if has_value {
            Ok(Value::Real(sum))
        } else {
            Ok(Value::Null)
        }
    }

    /// Compute AVG aggregate over window frame
    fn compute_aggregate_avg(
        partition: &[(usize, &Row)],
        pos: usize,
        expr: &Expr,
        frame: Option<&WindowFrame>,
        col_map: &HashMap<String, usize>,
    ) -> Result<Value> {
        let (start, end) = Self::compute_frame_bounds(pos, partition.len(), frame);
        let mut sum = 0.0;
        let mut count = 0;

        for (_, row) in partition.iter().skip(start).take(end - start) {
            let val = expr.eval(&row.values, col_map)?;
            if let Some(n) = val.as_f64() {
                sum += n;
                count += 1;
            }
        }

        if count > 0 {
            Ok(Value::Real(sum / count as f64))
        } else {
            Ok(Value::Null)
        }
    }

    /// Compute COUNT aggregate over window frame
    fn compute_aggregate_count(
        partition: &[(usize, &Row)],
        pos: usize,
        expr: Option<&Expr>,
        frame: Option<&WindowFrame>,
        col_map: &HashMap<String, usize>,
    ) -> Result<Value> {
        let (start, end) = Self::compute_frame_bounds(pos, partition.len(), frame);
        let mut count = 0i64;

        for (_, row) in partition.iter().skip(start).take(end - start) {
            match expr {
                Some(e) => {
                    let val = e.eval(&row.values, col_map)?;
                    if !val.is_null() {
                        count += 1;
                    }
                }
                None => {
                    // COUNT(*) - count all rows
                    count += 1;
                }
            }
        }

        Ok(Value::Integer(count))
    }

    /// Compute MIN aggregate over window frame
    fn compute_aggregate_min(
        partition: &[(usize, &Row)],
        pos: usize,
        expr: &Expr,
        frame: Option<&WindowFrame>,
        col_map: &HashMap<String, usize>,
    ) -> Result<Value> {
        let (start, end) = Self::compute_frame_bounds(pos, partition.len(), frame);
        let mut min: Option<Value> = None;

        for (_, row) in partition.iter().skip(start).take(end - start) {
            let val = expr.eval(&row.values, col_map)?;
            if !val.is_null() {
                min = Some(match min {
                    Some(m) if val < m => val,
                    Some(m) => m,
                    None => val,
                });
            }
        }

        Ok(min.unwrap_or(Value::Null))
    }

    /// Compute MAX aggregate over window frame
    fn compute_aggregate_max(
        partition: &[(usize, &Row)],
        pos: usize,
        expr: &Expr,
        frame: Option<&WindowFrame>,
        col_map: &HashMap<String, usize>,
    ) -> Result<Value> {
        let (start, end) = Self::compute_frame_bounds(pos, partition.len(), frame);
        let mut max: Option<Value> = None;

        for (_, row) in partition.iter().skip(start).take(end - start) {
            let val = expr.eval(&row.values, col_map)?;
            if !val.is_null() {
                max = Some(match max {
                    Some(m) if val > m => val,
                    Some(m) => m,
                    None => val,
                });
            }
        }

        Ok(max.unwrap_or(Value::Null))
    }
}

/// Build a column name to index map
fn build_col_map(columns: &[String]) -> HashMap<String, usize> {
    let mut map = HashMap::new();
    for (i, col) in columns.iter().enumerate() {
        map.insert(col.clone(), i);
        // Also add unqualified name
        if let Some(name) = col.split('.').next_back() {
            map.insert(name.to_string(), i);
        }
    }
    map
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{FrameBound, FrameUnit};

    fn make_row(values: Vec<Value>, columns: Vec<&str>) -> Row {
        Row::new(values, columns.into_iter().map(|s| s.to_string()).collect())
    }

    #[test]
    fn test_row_number() {
        let rows = vec![
            make_row(
                vec![Value::Integer(1), Value::Text("Alice".into())],
                vec!["id", "name"],
            ),
            make_row(
                vec![Value::Integer(2), Value::Text("Bob".into())],
                vec!["id", "name"],
            ),
            make_row(
                vec![Value::Integer(3), Value::Text("Carol".into())],
                vec!["id", "name"],
            ),
        ];

        let window_func = WindowFunction {
            function: WindowFunctionType::RowNumber,
            partition_by: vec![],
            order_by: vec![WindowOrderByExpr {
                expr: Expr::Column {
                    table: None,
                    name: "id".into(),
                    index: None,
                },
                order: WindowSortOrder::Asc,
                nulls_first: None,
            }],
            frame: None,
        };

        let window_exprs = vec![(Expr::Window(window_func), "row_num".to_string())];
        let result = WindowExecutor::execute(rows, &window_exprs).unwrap();

        assert_eq!(result.len(), 3);
        assert_eq!(result[0].values.last(), Some(&Value::Integer(1)));
        assert_eq!(result[1].values.last(), Some(&Value::Integer(2)));
        assert_eq!(result[2].values.last(), Some(&Value::Integer(3)));
    }

    #[test]
    fn test_partition_by() {
        let rows = vec![
            make_row(
                vec![Value::Text("A".into()), Value::Integer(10)],
                vec!["dept", "salary"],
            ),
            make_row(
                vec![Value::Text("A".into()), Value::Integer(20)],
                vec!["dept", "salary"],
            ),
            make_row(
                vec![Value::Text("B".into()), Value::Integer(15)],
                vec!["dept", "salary"],
            ),
            make_row(
                vec![Value::Text("B".into()), Value::Integer(25)],
                vec!["dept", "salary"],
            ),
        ];

        let window_func = WindowFunction {
            function: WindowFunctionType::RowNumber,
            partition_by: vec![Expr::Column {
                table: None,
                name: "dept".into(),
                index: None,
            }],
            order_by: vec![WindowOrderByExpr {
                expr: Expr::Column {
                    table: None,
                    name: "salary".into(),
                    index: None,
                },
                order: WindowSortOrder::Asc,
                nulls_first: None,
            }],
            frame: None,
        };

        let window_exprs = vec![(Expr::Window(window_func), "row_num".to_string())];
        let result = WindowExecutor::execute(rows, &window_exprs).unwrap();

        // Rows should be numbered within each partition
        // Department A: row 1 (salary 10), row 2 (salary 20)
        // Department B: row 1 (salary 15), row 2 (salary 25)
        assert_eq!(result.len(), 4);
    }

    #[test]
    fn test_running_sum() {
        let rows = vec![
            make_row(
                vec![Value::Integer(1), Value::Integer(100)],
                vec!["id", "amount"],
            ),
            make_row(
                vec![Value::Integer(2), Value::Integer(200)],
                vec!["id", "amount"],
            ),
            make_row(
                vec![Value::Integer(3), Value::Integer(150)],
                vec!["id", "amount"],
            ),
        ];

        let window_func = WindowFunction {
            function: WindowFunctionType::Sum(Box::new(Expr::Column {
                table: None,
                name: "amount".into(),
                index: None,
            })),
            partition_by: vec![],
            order_by: vec![WindowOrderByExpr {
                expr: Expr::Column {
                    table: None,
                    name: "id".into(),
                    index: None,
                },
                order: WindowSortOrder::Asc,
                nulls_first: None,
            }],
            frame: Some(WindowFrame {
                unit: FrameUnit::Rows,
                start: FrameBound::UnboundedPreceding,
                end: FrameBound::CurrentRow,
            }),
        };

        let window_exprs = vec![(Expr::Window(window_func), "running_total".to_string())];
        let result = WindowExecutor::execute(rows, &window_exprs).unwrap();

        assert_eq!(result.len(), 3);
        assert_eq!(result[0].values.last(), Some(&Value::Real(100.0)));
        assert_eq!(result[1].values.last(), Some(&Value::Real(300.0)));
        assert_eq!(result[2].values.last(), Some(&Value::Real(450.0)));
    }

    #[test]
    fn test_lag_lead() {
        let rows = vec![
            make_row(
                vec![Value::Integer(1), Value::Integer(100)],
                vec!["id", "value"],
            ),
            make_row(
                vec![Value::Integer(2), Value::Integer(200)],
                vec!["id", "value"],
            ),
            make_row(
                vec![Value::Integer(3), Value::Integer(300)],
                vec!["id", "value"],
            ),
        ];

        let lag_func = WindowFunction {
            function: WindowFunctionType::Lag {
                expr: Box::new(Expr::Column {
                    table: None,
                    name: "value".into(),
                    index: None,
                }),
                offset: 1,
                default: Some(Box::new(Expr::Literal(Value::Integer(0)))),
            },
            partition_by: vec![],
            order_by: vec![WindowOrderByExpr {
                expr: Expr::Column {
                    table: None,
                    name: "id".into(),
                    index: None,
                },
                order: WindowSortOrder::Asc,
                nulls_first: None,
            }],
            frame: None,
        };

        let window_exprs = vec![(Expr::Window(lag_func), "prev_value".to_string())];
        let result = WindowExecutor::execute(rows, &window_exprs).unwrap();

        assert_eq!(result.len(), 3);
        assert_eq!(result[0].values.last(), Some(&Value::Integer(0))); // Default for first row
        assert_eq!(result[1].values.last(), Some(&Value::Integer(100)));
        assert_eq!(result[2].values.last(), Some(&Value::Integer(200)));
    }
}