sql_cli/data/
arithmetic_evaluator.rs

1use crate::config::global::get_date_notation;
2use crate::data::data_view::DataView;
3use crate::data::datatable::{DataTable, DataValue};
4use crate::data::value_comparisons::compare_with_op;
5use crate::sql::aggregate_functions::AggregateFunctionRegistry; // New registry
6use crate::sql::aggregates::AggregateRegistry; // Old registry (for migration)
7use crate::sql::functions::FunctionRegistry;
8use crate::sql::parser::ast::{ColumnRef, WindowSpec};
9use crate::sql::recursive_parser::SqlExpression;
10use crate::sql::window_context::WindowContext;
11use crate::sql::window_functions::{ExpressionEvaluator, WindowFunctionRegistry};
12use anyhow::{anyhow, Result};
13use std::collections::{HashMap, HashSet};
14use std::sync::Arc;
15use std::time::Instant;
16use tracing::{debug, info};
17
18/// Evaluates SQL expressions to compute `DataValues` (for SELECT clauses)
19/// This is different from `RecursiveWhereEvaluator` which returns boolean
20pub struct ArithmeticEvaluator<'a> {
21    table: &'a DataTable,
22    _date_notation: String,
23    function_registry: Arc<FunctionRegistry>,
24    aggregate_registry: Arc<AggregateRegistry>, // Old registry (being phased out)
25    new_aggregate_registry: Arc<AggregateFunctionRegistry>, // New registry
26    window_function_registry: Arc<WindowFunctionRegistry>,
27    visible_rows: Option<Vec<usize>>, // For aggregate functions on filtered views
28    window_contexts: HashMap<u64, Arc<WindowContext>>, // Cache window contexts by hash
29    table_aliases: HashMap<String, String>, // Map alias -> table name for qualified columns
30}
31
32impl<'a> ArithmeticEvaluator<'a> {
33    #[must_use]
34    pub fn new(table: &'a DataTable) -> Self {
35        Self {
36            table,
37            _date_notation: get_date_notation(),
38            function_registry: Arc::new(FunctionRegistry::new()),
39            aggregate_registry: Arc::new(AggregateRegistry::new()),
40            new_aggregate_registry: Arc::new(AggregateFunctionRegistry::new()),
41            window_function_registry: Arc::new(WindowFunctionRegistry::new()),
42            visible_rows: None,
43            window_contexts: HashMap::new(),
44            table_aliases: HashMap::new(),
45        }
46    }
47
48    #[must_use]
49    pub fn with_date_notation(table: &'a DataTable, date_notation: String) -> Self {
50        Self {
51            table,
52            _date_notation: date_notation,
53            function_registry: Arc::new(FunctionRegistry::new()),
54            aggregate_registry: Arc::new(AggregateRegistry::new()),
55            new_aggregate_registry: Arc::new(AggregateFunctionRegistry::new()),
56            window_function_registry: Arc::new(WindowFunctionRegistry::new()),
57            visible_rows: None,
58            window_contexts: HashMap::new(),
59            table_aliases: HashMap::new(),
60        }
61    }
62
63    /// Set visible rows for aggregate functions (for filtered views)
64    #[must_use]
65    pub fn with_visible_rows(mut self, rows: Vec<usize>) -> Self {
66        self.visible_rows = Some(rows);
67        self
68    }
69
70    /// Set table aliases for qualified column resolution
71    #[must_use]
72    pub fn with_table_aliases(mut self, aliases: HashMap<String, String>) -> Self {
73        self.table_aliases = aliases;
74        self
75    }
76
77    #[must_use]
78    pub fn with_date_notation_and_registry(
79        table: &'a DataTable,
80        date_notation: String,
81        function_registry: Arc<FunctionRegistry>,
82    ) -> Self {
83        Self {
84            table,
85            _date_notation: date_notation,
86            function_registry,
87            aggregate_registry: Arc::new(AggregateRegistry::new()),
88            new_aggregate_registry: Arc::new(AggregateFunctionRegistry::new()),
89            window_function_registry: Arc::new(WindowFunctionRegistry::new()),
90            visible_rows: None,
91            window_contexts: HashMap::new(),
92            table_aliases: HashMap::new(),
93        }
94    }
95
96    /// Find a column name similar to the given name using edit distance
97    fn find_similar_column(&self, name: &str) -> Option<String> {
98        let columns = self.table.column_names();
99        let mut best_match: Option<(String, usize)> = None;
100
101        for col in columns {
102            let distance = self.edit_distance(&col.to_lowercase(), &name.to_lowercase());
103            // Only suggest if distance is small (likely a typo)
104            // Allow up to 3 edits for longer names
105            let max_distance = if name.len() > 10 { 3 } else { 2 };
106            if distance <= max_distance {
107                match &best_match {
108                    None => best_match = Some((col, distance)),
109                    Some((_, best_dist)) if distance < *best_dist => {
110                        best_match = Some((col, distance));
111                    }
112                    _ => {}
113                }
114            }
115        }
116
117        best_match.map(|(name, _)| name)
118    }
119
120    /// Calculate Levenshtein edit distance between two strings
121    fn edit_distance(&self, s1: &str, s2: &str) -> usize {
122        // Use the shared implementation from string_methods
123        crate::sql::functions::string_methods::EditDistanceFunction::calculate_edit_distance(s1, s2)
124    }
125
126    /// Evaluate an SQL expression to produce a `DataValue`
127    pub fn evaluate(&mut self, expr: &SqlExpression, row_index: usize) -> Result<DataValue> {
128        debug!(
129            "ArithmeticEvaluator: evaluating {:?} for row {}",
130            expr, row_index
131        );
132
133        match expr {
134            SqlExpression::Column(column_ref) => self.evaluate_column_ref(column_ref, row_index),
135            SqlExpression::StringLiteral(s) => Ok(DataValue::String(s.clone())),
136            SqlExpression::BooleanLiteral(b) => Ok(DataValue::Boolean(*b)),
137            SqlExpression::NumberLiteral(n) => self.evaluate_number_literal(n),
138            SqlExpression::Null => Ok(DataValue::Null),
139            SqlExpression::BinaryOp { left, op, right } => {
140                self.evaluate_binary_op(left, op, right, row_index)
141            }
142            SqlExpression::FunctionCall {
143                name,
144                args,
145                distinct,
146            } => self.evaluate_function_with_distinct(name, args, *distinct, row_index),
147            SqlExpression::WindowFunction {
148                name,
149                args,
150                window_spec,
151            } => self.evaluate_window_function(name, args, window_spec, row_index),
152            SqlExpression::MethodCall {
153                object,
154                method,
155                args,
156            } => self.evaluate_method_call(object, method, args, row_index),
157            SqlExpression::ChainedMethodCall { base, method, args } => {
158                // Evaluate the base expression first, then apply the method
159                let base_value = self.evaluate(base, row_index)?;
160                self.evaluate_method_on_value(&base_value, method, args, row_index)
161            }
162            SqlExpression::Between { expr, lower, upper } => {
163                let val = self.evaluate(expr, row_index)?;
164                let lo = self.evaluate(lower, row_index)?;
165                let hi = self.evaluate(upper, row_index)?;
166                let ge = compare_with_op(&val, &lo, ">=", false);
167                let le = compare_with_op(&val, &hi, "<=", false);
168                Ok(DataValue::Boolean(ge && le))
169            }
170            // Logical negation as a value-producing expression. Reached, e.g., by
171            // a post-aggregation `HAVING NOT (COUNT(*) > 2)` predicate (P10).
172            // Three-valued logic: NOT NULL is NULL; otherwise negate the boolean.
173            SqlExpression::Not { expr } => {
174                let inner = self.evaluate(expr, row_index)?;
175                match inner {
176                    DataValue::Null => Ok(DataValue::Null),
177                    other => Ok(DataValue::Boolean(!self.to_bool(&other)?)),
178                }
179            }
180            // IN / NOT IN as a value-producing expression — needed when they
181            // appear inside CASE branches or other arithmetic contexts. The
182            // WHERE path has its own evaluator; this mirrors its equality
183            // semantics via compare_with_op. After subquery rewriting, an
184            // `x IN (SELECT ...)` arrives here as an InList of literals.
185            SqlExpression::InList { expr, values } => {
186                let val = self.evaluate(expr, row_index)?;
187                for v in values {
188                    let item = self.evaluate(v, row_index)?;
189                    if compare_with_op(&val, &item, "=", false) {
190                        return Ok(DataValue::Boolean(true));
191                    }
192                }
193                Ok(DataValue::Boolean(false))
194            }
195            SqlExpression::NotInList { expr, values } => {
196                let val = self.evaluate(expr, row_index)?;
197                for v in values {
198                    let item = self.evaluate(v, row_index)?;
199                    if compare_with_op(&val, &item, "=", false) {
200                        return Ok(DataValue::Boolean(false));
201                    }
202                }
203                Ok(DataValue::Boolean(true))
204            }
205            SqlExpression::CaseExpression {
206                when_branches,
207                else_branch,
208            } => self.evaluate_case_expression(when_branches, else_branch, row_index),
209            SqlExpression::SimpleCaseExpression {
210                expr,
211                when_branches,
212                else_branch,
213            } => self.evaluate_simple_case_expression(expr, when_branches, else_branch, row_index),
214            SqlExpression::DateTimeConstructor {
215                year,
216                month,
217                day,
218                hour,
219                minute,
220                second,
221            } => self.evaluate_datetime_constructor(*year, *month, *day, *hour, *minute, *second),
222            SqlExpression::DateTimeToday {
223                hour,
224                minute,
225                second,
226            } => self.evaluate_datetime_today(*hour, *minute, *second),
227            _ => Err(anyhow!(
228                "Unsupported expression type for arithmetic evaluation: {:?}",
229                expr
230            )),
231        }
232    }
233
234    /// Evaluate a column reference with proper table scoping
235    fn evaluate_column_ref(&self, column_ref: &ColumnRef, row_index: usize) -> Result<DataValue> {
236        if let Some(table_prefix) = &column_ref.table_prefix {
237            // Resolve alias if it exists in table_aliases map
238            let actual_table = self
239                .table_aliases
240                .get(table_prefix)
241                .map(|s| s.as_str())
242                .unwrap_or(table_prefix);
243
244            // Try qualified lookup with resolved table name
245            let qualified_name = format!("{}.{}", actual_table, column_ref.name);
246
247            if let Some(col_idx) = self.table.find_column_by_qualified_name(&qualified_name) {
248                debug!(
249                    "Resolved {}.{} -> '{}' at index {}",
250                    table_prefix, column_ref.name, qualified_name, col_idx
251                );
252                return self
253                    .table
254                    .get_value(row_index, col_idx)
255                    .ok_or_else(|| anyhow!("Row {} out of bounds", row_index))
256                    .map(|v| v.clone());
257            }
258
259            // Fallback: try unqualified lookup
260            if let Some(col_idx) = self.table.get_column_index(&column_ref.name) {
261                debug!(
262                    "Resolved {}.{} -> unqualified '{}' at index {}",
263                    table_prefix, column_ref.name, column_ref.name, col_idx
264                );
265                return self
266                    .table
267                    .get_value(row_index, col_idx)
268                    .ok_or_else(|| anyhow!("Row {} out of bounds", row_index))
269                    .map(|v| v.clone());
270            }
271
272            // If not found, return error
273            Err(anyhow!(
274                "Column '{}' not found. Table '{}' may not support qualified column names",
275                qualified_name,
276                actual_table
277            ))
278        } else {
279            // Simple column name lookup
280            self.evaluate_column(&column_ref.name, row_index)
281        }
282    }
283
284    /// Evaluate a column reference
285    fn evaluate_column(&self, column_name: &str, row_index: usize) -> Result<DataValue> {
286        // First try to resolve qualified column names (table.column or alias.column)
287        let resolved_column = if column_name.contains('.') {
288            // Split on last dot to handle cases like "schema.table.column"
289            if let Some(dot_pos) = column_name.rfind('.') {
290                let _table_or_alias = &column_name[..dot_pos];
291                let col_name = &column_name[dot_pos + 1..];
292
293                // For now, just use the column name part
294                // In the future, we could validate the table/alias part
295                debug!(
296                    "Resolving qualified column: {} -> {}",
297                    column_name, col_name
298                );
299                col_name.to_string()
300            } else {
301                column_name.to_string()
302            }
303        } else {
304            column_name.to_string()
305        };
306
307        let col_index = if let Some(idx) = self.table.get_column_index(&resolved_column) {
308            idx
309        } else if resolved_column != column_name {
310            // If not found, try the original name
311            if let Some(idx) = self.table.get_column_index(column_name) {
312                idx
313            } else {
314                let suggestion = self.find_similar_column(&resolved_column);
315                return Err(match suggestion {
316                    Some(similar) => anyhow!(
317                        "Column '{}' not found. Did you mean '{}'?",
318                        column_name,
319                        similar
320                    ),
321                    None => anyhow!("Column '{}' not found", column_name),
322                });
323            }
324        } else {
325            let suggestion = self.find_similar_column(&resolved_column);
326            return Err(match suggestion {
327                Some(similar) => anyhow!(
328                    "Column '{}' not found. Did you mean '{}'?",
329                    column_name,
330                    similar
331                ),
332                None => anyhow!("Column '{}' not found", column_name),
333            });
334        };
335
336        if row_index >= self.table.row_count() {
337            return Err(anyhow!("Row index {} out of bounds", row_index));
338        }
339
340        let row = self
341            .table
342            .get_row(row_index)
343            .ok_or_else(|| anyhow!("Row {} not found", row_index))?;
344
345        let value = row
346            .get(col_index)
347            .ok_or_else(|| anyhow!("Column index {} out of bounds for row", col_index))?;
348
349        Ok(value.clone())
350    }
351
352    /// Evaluate a number literal (handles both integers and floats)
353    fn evaluate_number_literal(&self, number_str: &str) -> Result<DataValue> {
354        // Try to parse as integer first
355        if let Ok(int_val) = number_str.parse::<i64>() {
356            return Ok(DataValue::Integer(int_val));
357        }
358
359        // If that fails, try as float
360        if let Ok(float_val) = number_str.parse::<f64>() {
361            return Ok(DataValue::Float(float_val));
362        }
363
364        Err(anyhow!("Invalid number literal: {}", number_str))
365    }
366
367    /// Evaluate a binary operation (arithmetic)
368    fn evaluate_binary_op(
369        &mut self,
370        left: &SqlExpression,
371        op: &str,
372        right: &SqlExpression,
373        row_index: usize,
374    ) -> Result<DataValue> {
375        let left_val = self.evaluate(left, row_index)?;
376        let right_val = self.evaluate(right, row_index)?;
377
378        debug!(
379            "ArithmeticEvaluator: {} {} {}",
380            self.format_value(&left_val),
381            op,
382            self.format_value(&right_val)
383        );
384
385        match op {
386            "+" => self.add_values(&left_val, &right_val),
387            "-" => self.subtract_values(&left_val, &right_val),
388            "*" => self.multiply_values(&left_val, &right_val),
389            "/" => self.divide_values(&left_val, &right_val),
390            "%" => {
391                // Modulo operator - call MOD function
392                let args = vec![left.clone(), right.clone()];
393                self.evaluate_function("MOD", &args, row_index)
394            }
395            // Comparison operators (return boolean results)
396            // Use centralized comparison logic for consistency
397            ">" | "<" | ">=" | "<=" | "=" | "!=" | "<>" => {
398                let result = compare_with_op(&left_val, &right_val, op, false);
399                Ok(DataValue::Boolean(result))
400            }
401            // IS NULL / IS NOT NULL operators
402            "IS NULL" => Ok(DataValue::Boolean(matches!(left_val, DataValue::Null))),
403            "IS NOT NULL" => Ok(DataValue::Boolean(!matches!(left_val, DataValue::Null))),
404            // Logical operators
405            "AND" => {
406                let left_bool = self.to_bool(&left_val)?;
407                let right_bool = self.to_bool(&right_val)?;
408                Ok(DataValue::Boolean(left_bool && right_bool))
409            }
410            "OR" => {
411                let left_bool = self.to_bool(&left_val)?;
412                let right_bool = self.to_bool(&right_val)?;
413                Ok(DataValue::Boolean(left_bool || right_bool))
414            }
415            // LIKE operator - SQL pattern matching
416            "LIKE" => {
417                let text = self.value_to_string(&left_val);
418                let pattern = self.value_to_string(&right_val);
419                let matches = self.sql_like_match(&text, &pattern);
420                Ok(DataValue::Boolean(matches))
421            }
422            _ => Err(anyhow!("Unsupported arithmetic operator: {}", op)),
423        }
424    }
425
426    /// Add two `DataValues` with type coercion
427    fn add_values(&self, left: &DataValue, right: &DataValue) -> Result<DataValue> {
428        // NULL handling - any operation with NULL returns NULL
429        if matches!(left, DataValue::Null) || matches!(right, DataValue::Null) {
430            return Ok(DataValue::Null);
431        }
432
433        match (left, right) {
434            (DataValue::Integer(a), DataValue::Integer(b)) => Ok(DataValue::Integer(a + b)),
435            (DataValue::Integer(a), DataValue::Float(b)) => Ok(DataValue::Float(*a as f64 + b)),
436            (DataValue::Float(a), DataValue::Integer(b)) => Ok(DataValue::Float(a + *b as f64)),
437            (DataValue::Float(a), DataValue::Float(b)) => Ok(DataValue::Float(a + b)),
438            _ => Err(anyhow!("Cannot add {:?} and {:?}", left, right)),
439        }
440    }
441
442    /// Subtract two `DataValues` with type coercion
443    fn subtract_values(&self, left: &DataValue, right: &DataValue) -> Result<DataValue> {
444        // NULL handling - any operation with NULL returns NULL
445        if matches!(left, DataValue::Null) || matches!(right, DataValue::Null) {
446            return Ok(DataValue::Null);
447        }
448
449        match (left, right) {
450            (DataValue::Integer(a), DataValue::Integer(b)) => Ok(DataValue::Integer(a - b)),
451            (DataValue::Integer(a), DataValue::Float(b)) => Ok(DataValue::Float(*a as f64 - b)),
452            (DataValue::Float(a), DataValue::Integer(b)) => Ok(DataValue::Float(a - *b as f64)),
453            (DataValue::Float(a), DataValue::Float(b)) => Ok(DataValue::Float(a - b)),
454            _ => Err(anyhow!("Cannot subtract {:?} and {:?}", left, right)),
455        }
456    }
457
458    /// Multiply two `DataValues` with type coercion
459    fn multiply_values(&self, left: &DataValue, right: &DataValue) -> Result<DataValue> {
460        // NULL handling - any operation with NULL returns NULL
461        if matches!(left, DataValue::Null) || matches!(right, DataValue::Null) {
462            return Ok(DataValue::Null);
463        }
464
465        match (left, right) {
466            (DataValue::Integer(a), DataValue::Integer(b)) => Ok(DataValue::Integer(a * b)),
467            (DataValue::Integer(a), DataValue::Float(b)) => Ok(DataValue::Float(*a as f64 * b)),
468            (DataValue::Float(a), DataValue::Integer(b)) => Ok(DataValue::Float(a * *b as f64)),
469            (DataValue::Float(a), DataValue::Float(b)) => Ok(DataValue::Float(a * b)),
470            _ => Err(anyhow!("Cannot multiply {:?} and {:?}", left, right)),
471        }
472    }
473
474    /// Divide two `DataValues` with type coercion
475    fn divide_values(&self, left: &DataValue, right: &DataValue) -> Result<DataValue> {
476        // NULL handling - any operation with NULL returns NULL
477        if matches!(left, DataValue::Null) || matches!(right, DataValue::Null) {
478            return Ok(DataValue::Null);
479        }
480
481        // Check for division by zero first
482        let is_zero = match right {
483            DataValue::Integer(0) => true,
484            DataValue::Float(f) if *f == 0.0 => true, // Only check for exact zero, not epsilon
485            _ => false,
486        };
487
488        if is_zero {
489            return Err(anyhow!("Division by zero"));
490        }
491
492        match (left, right) {
493            (DataValue::Integer(a), DataValue::Integer(b)) => {
494                // Integer division - if result is exact, keep as int, otherwise promote to float
495                if a % b == 0 {
496                    Ok(DataValue::Integer(a / b))
497                } else {
498                    Ok(DataValue::Float(*a as f64 / *b as f64))
499                }
500            }
501            (DataValue::Integer(a), DataValue::Float(b)) => Ok(DataValue::Float(*a as f64 / b)),
502            (DataValue::Float(a), DataValue::Integer(b)) => Ok(DataValue::Float(a / *b as f64)),
503            (DataValue::Float(a), DataValue::Float(b)) => Ok(DataValue::Float(a / b)),
504            _ => Err(anyhow!("Cannot divide {:?} and {:?}", left, right)),
505        }
506    }
507
508    /// Format a `DataValue` for debug output
509    fn format_value(&self, value: &DataValue) -> String {
510        match value {
511            DataValue::Integer(i) => i.to_string(),
512            DataValue::Float(f) => f.to_string(),
513            DataValue::String(s) => format!("'{s}'"),
514            _ => format!("{value:?}"),
515        }
516    }
517
518    /// Convert a DataValue to boolean for logical operations
519    fn to_bool(&self, value: &DataValue) -> Result<bool> {
520        match value {
521            DataValue::Boolean(b) => Ok(*b),
522            DataValue::Integer(i) => Ok(*i != 0),
523            DataValue::Float(f) => Ok(*f != 0.0),
524            DataValue::Null => Ok(false),
525            _ => Err(anyhow!("Cannot convert {:?} to boolean", value)),
526        }
527    }
528
529    /// Convert DataValue to string for pattern matching
530    fn value_to_string(&self, value: &DataValue) -> String {
531        match value {
532            DataValue::String(s) => s.clone(),
533            DataValue::InternedString(s) => s.to_string(),
534            DataValue::Integer(i) => i.to_string(),
535            DataValue::Float(f) => f.to_string(),
536            DataValue::Boolean(b) => b.to_string(),
537            DataValue::DateTime(dt) => dt.to_string(),
538            DataValue::Vector(v) => {
539                // Format as "[x,y,z]"
540                let components: Vec<String> = v.iter().map(|f| f.to_string()).collect();
541                format!("[{}]", components.join(","))
542            }
543            DataValue::Null => String::new(),
544        }
545    }
546
547    /// SQL LIKE pattern matching
548    /// Supports % (any chars) and _ (single char)
549    fn sql_like_match(&self, text: &str, pattern: &str) -> bool {
550        let pattern_chars: Vec<char> = pattern.chars().collect();
551        let text_chars: Vec<char> = text.chars().collect();
552
553        self.like_match_recursive(&text_chars, 0, &pattern_chars, 0)
554    }
555
556    /// Recursive helper for LIKE matching
557    fn like_match_recursive(
558        &self,
559        text: &[char],
560        text_pos: usize,
561        pattern: &[char],
562        pattern_pos: usize,
563    ) -> bool {
564        // If we've consumed both text and pattern, it's a match
565        if pattern_pos >= pattern.len() {
566            return text_pos >= text.len();
567        }
568
569        // Handle % wildcard (matches zero or more characters)
570        if pattern[pattern_pos] == '%' {
571            // Try matching zero characters (skip the %)
572            if self.like_match_recursive(text, text_pos, pattern, pattern_pos + 1) {
573                return true;
574            }
575            // Try matching one or more characters
576            if text_pos < text.len() {
577                return self.like_match_recursive(text, text_pos + 1, pattern, pattern_pos);
578            }
579            return false;
580        }
581
582        // If text is consumed but pattern isn't, no match
583        if text_pos >= text.len() {
584            return false;
585        }
586
587        // Handle _ wildcard (matches exactly one character)
588        if pattern[pattern_pos] == '_' {
589            return self.like_match_recursive(text, text_pos + 1, pattern, pattern_pos + 1);
590        }
591
592        // Handle literal character match
593        if text[text_pos] == pattern[pattern_pos] {
594            return self.like_match_recursive(text, text_pos + 1, pattern, pattern_pos + 1);
595        }
596
597        false
598    }
599
600    /// Evaluate a function call
601    fn evaluate_function_with_distinct(
602        &mut self,
603        name: &str,
604        args: &[SqlExpression],
605        distinct: bool,
606        row_index: usize,
607    ) -> Result<DataValue> {
608        // If DISTINCT is specified, handle it specially for aggregate functions
609        if distinct {
610            let name_upper = name.to_uppercase();
611
612            // Check if it's an aggregate function in either registry
613            if self.aggregate_registry.is_aggregate(&name_upper)
614                || self.new_aggregate_registry.contains(&name_upper)
615            {
616                return self.evaluate_aggregate_with_distinct(&name_upper, args, row_index);
617            } else {
618                return Err(anyhow!(
619                    "DISTINCT can only be used with aggregate functions"
620                ));
621            }
622        }
623
624        // Otherwise, use the regular evaluation
625        self.evaluate_function(name, args, row_index)
626    }
627
628    fn evaluate_aggregate_with_distinct(
629        &mut self,
630        name: &str,
631        args: &[SqlExpression],
632        _row_index: usize,
633    ) -> Result<DataValue> {
634        let name_upper = name.to_uppercase();
635
636        // Check new aggregate registry first for migrated functions
637        if self.new_aggregate_registry.get(&name_upper).is_some() {
638            let rows_to_process: Vec<usize> = if let Some(ref visible) = self.visible_rows {
639                visible.clone()
640            } else {
641                (0..self.table.rows.len()).collect()
642            };
643
644            // Collect and deduplicate values for DISTINCT
645            let mut vals = Vec::new();
646            for &row_idx in &rows_to_process {
647                if !args.is_empty() {
648                    let value = self.evaluate(&args[0], row_idx)?;
649                    vals.push(value);
650                }
651            }
652
653            // Deduplicate values
654            let mut seen = HashSet::new();
655            let unique_values: Vec<_> = vals
656                .into_iter()
657                .filter(|v| {
658                    let key = format!("{:?}", v);
659                    seen.insert(key)
660                })
661                .collect();
662
663            // Get the aggregate function from the new registry
664            let agg_func = self.new_aggregate_registry.get(&name_upper).unwrap();
665            let mut state = agg_func.create_state();
666
667            // Use unique values
668            for value in &unique_values {
669                state.accumulate(value)?;
670            }
671
672            return Ok(state.finalize());
673        }
674
675        // Check old aggregate registry (DISTINCT handling)
676        if self.aggregate_registry.get(&name_upper).is_some() {
677            // Determine which rows to process first
678            let rows_to_process: Vec<usize> = if let Some(ref visible) = self.visible_rows {
679                visible.clone()
680            } else {
681                (0..self.table.rows.len()).collect()
682            };
683
684            // Special handling for STRING_AGG with separator parameter
685            if name_upper == "STRING_AGG" && args.len() >= 2 {
686                // STRING_AGG(DISTINCT column, separator)
687                let mut state = crate::sql::aggregates::AggregateState::StringAgg(
688                    // Evaluate the separator (second argument) once
689                    if args.len() >= 2 {
690                        let separator = self.evaluate(&args[1], 0)?; // Separator doesn't depend on row
691                        match separator {
692                            DataValue::String(s) => crate::sql::aggregates::StringAggState::new(&s),
693                            DataValue::InternedString(s) => {
694                                crate::sql::aggregates::StringAggState::new(&s)
695                            }
696                            _ => crate::sql::aggregates::StringAggState::new(","), // Default separator
697                        }
698                    } else {
699                        crate::sql::aggregates::StringAggState::new(",")
700                    },
701                );
702
703                // Evaluate the first argument (column) for each row and accumulate
704                // Handle DISTINCT - use a HashSet to track seen values
705                let mut seen_values = HashSet::new();
706
707                for &row_idx in &rows_to_process {
708                    let value = self.evaluate(&args[0], row_idx)?;
709
710                    // Skip if we've seen this value
711                    if !seen_values.insert(value.clone()) {
712                        continue; // Skip duplicate values
713                    }
714
715                    // Now get the aggregate function and accumulate
716                    let agg_func = self.aggregate_registry.get(&name_upper).unwrap();
717                    agg_func.accumulate(&mut state, &value)?;
718                }
719
720                // Finalize the aggregate
721                let agg_func = self.aggregate_registry.get(&name_upper).unwrap();
722                return Ok(agg_func.finalize(state));
723            }
724
725            // For other aggregates with DISTINCT
726            // Evaluate the argument expression for each row
727            let mut vals = Vec::new();
728            for &row_idx in &rows_to_process {
729                if !args.is_empty() {
730                    let value = self.evaluate(&args[0], row_idx)?;
731                    vals.push(value);
732                }
733            }
734
735            // Deduplicate values for DISTINCT
736            let mut seen = HashSet::new();
737            let mut unique_values = Vec::new();
738            for value in vals {
739                if seen.insert(value.clone()) {
740                    unique_values.push(value);
741                }
742            }
743
744            // Now get the aggregate function and process
745            let agg_func = self.aggregate_registry.get(&name_upper).unwrap();
746            let mut state = agg_func.init();
747
748            // Use unique values
749            for value in &unique_values {
750                agg_func.accumulate(&mut state, value)?;
751            }
752
753            return Ok(agg_func.finalize(state));
754        }
755
756        Err(anyhow!("Unknown aggregate function: {}", name))
757    }
758
759    fn evaluate_function(
760        &mut self,
761        name: &str,
762        args: &[SqlExpression],
763        row_index: usize,
764    ) -> Result<DataValue> {
765        // Check if this is an aggregate function
766        let name_upper = name.to_uppercase();
767
768        // Check new aggregate registry first (for migrated functions)
769        if self.new_aggregate_registry.get(&name_upper).is_some() {
770            // Use new registry for SUM
771            let rows_to_process: Vec<usize> = if let Some(ref visible) = self.visible_rows {
772                visible.clone()
773            } else {
774                (0..self.table.rows.len()).collect()
775            };
776
777            // Get the aggregate function from the new registry
778            let agg_func = self.new_aggregate_registry.get(&name_upper).unwrap();
779            let mut state = agg_func.create_state();
780
781            // Special handling for COUNT(*)
782            if name_upper == "COUNT" || name_upper == "COUNT_STAR" {
783                if args.is_empty()
784                    || (args.len() == 1
785                        && matches!(&args[0], SqlExpression::Column(col) if col.name == "*"))
786                    || (args.len() == 1
787                        && matches!(&args[0], SqlExpression::StringLiteral(s) if s == "*"))
788                {
789                    // COUNT(*) or COUNT_STAR - count all rows
790                    for _ in &rows_to_process {
791                        state.accumulate(&DataValue::Integer(1))?;
792                    }
793                } else {
794                    // COUNT(column) - count non-null values
795                    for &row_idx in &rows_to_process {
796                        let value = self.evaluate(&args[0], row_idx)?;
797                        state.accumulate(&value)?;
798                    }
799                }
800            } else {
801                // Other aggregates - evaluate arguments and accumulate
802                if !args.is_empty() {
803                    for &row_idx in &rows_to_process {
804                        let value = self.evaluate(&args[0], row_idx)?;
805                        state.accumulate(&value)?;
806                    }
807                }
808            }
809
810            return Ok(state.finalize());
811        }
812
813        // Check old aggregate registry (for non-migrated functions)
814        if self.aggregate_registry.get(&name_upper).is_some() {
815            // Determine which rows to process first
816            let rows_to_process: Vec<usize> = if let Some(ref visible) = self.visible_rows {
817                visible.clone()
818            } else {
819                (0..self.table.rows.len()).collect()
820            };
821
822            // Special handling for STRING_AGG with separator parameter
823            if name_upper == "STRING_AGG" && args.len() >= 2 {
824                // STRING_AGG(column, separator) - without DISTINCT (handled separately)
825                let mut state = crate::sql::aggregates::AggregateState::StringAgg(
826                    // Evaluate the separator (second argument) once
827                    if args.len() >= 2 {
828                        let separator = self.evaluate(&args[1], 0)?; // Separator doesn't depend on row
829                        match separator {
830                            DataValue::String(s) => crate::sql::aggregates::StringAggState::new(&s),
831                            DataValue::InternedString(s) => {
832                                crate::sql::aggregates::StringAggState::new(&s)
833                            }
834                            _ => crate::sql::aggregates::StringAggState::new(","), // Default separator
835                        }
836                    } else {
837                        crate::sql::aggregates::StringAggState::new(",")
838                    },
839                );
840
841                // Evaluate the first argument (column) for each row and accumulate
842                for &row_idx in &rows_to_process {
843                    let value = self.evaluate(&args[0], row_idx)?;
844                    // Now get the aggregate function and accumulate
845                    let agg_func = self.aggregate_registry.get(&name_upper).unwrap();
846                    agg_func.accumulate(&mut state, &value)?;
847                }
848
849                // Finalize the aggregate
850                let agg_func = self.aggregate_registry.get(&name_upper).unwrap();
851                return Ok(agg_func.finalize(state));
852            }
853
854            // Evaluate arguments first if needed (to avoid borrow issues)
855            let values = if !args.is_empty()
856                && !(args.len() == 1
857                    && matches!(&args[0], SqlExpression::Column(c) if c.name == "*"))
858            {
859                // Evaluate the argument expression for each row
860                let mut vals = Vec::new();
861                for &row_idx in &rows_to_process {
862                    let value = self.evaluate(&args[0], row_idx)?;
863                    vals.push(value);
864                }
865                Some(vals)
866            } else {
867                None
868            };
869
870            // Now get the aggregate function and process
871            let agg_func = self.aggregate_registry.get(&name_upper).unwrap();
872            let mut state = agg_func.init();
873
874            if let Some(values) = values {
875                // Use evaluated values (DISTINCT is handled in evaluate_aggregate_with_distinct)
876                for value in &values {
877                    agg_func.accumulate(&mut state, value)?;
878                }
879            } else {
880                // COUNT(*) case
881                for _ in &rows_to_process {
882                    agg_func.accumulate(&mut state, &DataValue::Integer(1))?;
883                }
884            }
885
886            return Ok(agg_func.finalize(state));
887        }
888
889        // First check if this function exists in the registry
890        if self.function_registry.get(name).is_some() {
891            // Evaluate all arguments first to avoid borrow issues
892            let mut evaluated_args = Vec::new();
893            for arg in args {
894                evaluated_args.push(self.evaluate(arg, row_index)?);
895            }
896
897            // Get the function and call it
898            let func = self.function_registry.get(name).unwrap();
899            return func.evaluate(&evaluated_args);
900        }
901
902        // If not in registry, return error for unknown function
903        Err(anyhow!("Unknown function: {}", name))
904    }
905
906    /// Get or create a WindowContext for the given specification
907    /// Public to allow pre-creation of contexts in query engine (optimization)
908    pub fn get_or_create_window_context(
909        &mut self,
910        spec: &WindowSpec,
911    ) -> Result<Arc<WindowContext>> {
912        let overall_start = Instant::now();
913
914        // Create a hash-based key for fast caching (much faster than format!("{:?}", spec))
915        let key = spec.compute_hash();
916
917        if let Some(context) = self.window_contexts.get(&key) {
918            info!(
919                "WindowContext cache hit for spec (lookup: {:.2}μs)",
920                overall_start.elapsed().as_micros()
921            );
922            return Ok(Arc::clone(context));
923        }
924
925        info!("WindowContext cache miss - creating new context");
926        let dataview_start = Instant::now();
927
928        // Create a DataView from the table, restricted to the visible rows when the
929        // query filtered. Window functions must partition over the post-WHERE row set:
930        // SQL evaluates them after FROM/WHERE/GROUP BY/HAVING, so a filtered-out row
931        // must not appear in a partition, occupy a ROW_NUMBER slot, or be counted.
932        //
933        // The indices here are source-table indices, which is what DataView::with_rows
934        // expects and what WindowContext reads back via get_visible_rows() - so the
935        // whole path stays in one index space.
936        let data_view = if let Some(ref visible_rows) = self.visible_rows {
937            DataView::new(Arc::new(self.table.clone())).with_rows(visible_rows.clone())
938        } else {
939            DataView::new(Arc::new(self.table.clone()))
940        };
941
942        info!(
943            "DataView creation took {:.2}μs",
944            dataview_start.elapsed().as_micros()
945        );
946        let context_start = Instant::now();
947
948        // Create the WindowContext with the full spec (including frame)
949        let context = WindowContext::new_with_spec(Arc::new(data_view), spec.clone())?;
950
951        info!(
952            "WindowContext::new_with_spec took {:.2}ms (rows: {})",
953            context_start.elapsed().as_secs_f64() * 1000.0,
954            self.table.row_count()
955        );
956
957        let context = Arc::new(context);
958        self.window_contexts.insert(key, Arc::clone(&context));
959
960        info!(
961            "Total WindowContext creation (cache miss) took {:.2}ms",
962            overall_start.elapsed().as_secs_f64() * 1000.0
963        );
964
965        Ok(context)
966    }
967
968    /// Evaluate a window function
969    fn evaluate_window_function(
970        &mut self,
971        name: &str,
972        args: &[SqlExpression],
973        spec: &WindowSpec,
974        row_index: usize,
975    ) -> Result<DataValue> {
976        let func_start = Instant::now();
977        let name_upper = name.to_uppercase();
978
979        // First check if this is a syntactic sugar function in the registry
980        debug!("Looking for window function {} in registry", name_upper);
981        if let Some(window_fn_arc) = self.window_function_registry.get(&name_upper) {
982            debug!("Found window function {} in registry", name_upper);
983
984            // Dereference to get the actual window function
985            let window_fn = window_fn_arc.as_ref();
986
987            // Validate arguments
988            window_fn.validate_args(args)?;
989
990            // Transform the window spec based on the function's requirements
991            let transformed_spec = window_fn.transform_window_spec(spec, args)?;
992
993            // Get or create the window context with the transformed spec
994            let context = self.get_or_create_window_context(&transformed_spec)?;
995
996            // Create an expression evaluator adapter
997            struct EvaluatorAdapter<'a, 'b> {
998                evaluator: &'a mut ArithmeticEvaluator<'b>,
999                row_index: usize,
1000            }
1001
1002            impl<'a, 'b> ExpressionEvaluator for EvaluatorAdapter<'a, 'b> {
1003                fn evaluate(
1004                    &mut self,
1005                    expr: &SqlExpression,
1006                    row_index: usize,
1007                ) -> Result<DataValue> {
1008                    self.evaluator.evaluate(expr, row_index)
1009                }
1010            }
1011
1012            let mut adapter = EvaluatorAdapter {
1013                evaluator: self,
1014                row_index,
1015            };
1016
1017            let compute_start = Instant::now();
1018            // Call the window function's compute method
1019            let result = window_fn.compute(&context, row_index, args, &mut adapter);
1020
1021            info!(
1022                "{} (registry) evaluation: total={:.2}μs, compute={:.2}μs",
1023                name_upper,
1024                func_start.elapsed().as_micros(),
1025                compute_start.elapsed().as_micros()
1026            );
1027
1028            return result;
1029        }
1030
1031        // Fall back to built-in window functions
1032        let context_start = Instant::now();
1033        let context = self.get_or_create_window_context(spec)?;
1034        let context_time = context_start.elapsed();
1035
1036        let eval_start = Instant::now();
1037
1038        let result = match name_upper.as_str() {
1039            "LAG" => {
1040                // LAG(column, offset, default)
1041                if args.is_empty() {
1042                    return Err(anyhow!("LAG requires at least 1 argument"));
1043                }
1044
1045                // Get column name
1046                let column = match &args[0] {
1047                    SqlExpression::Column(col) => col.clone(),
1048                    _ => return Err(anyhow!("LAG first argument must be a column")),
1049                };
1050
1051                // Get offset (default 1)
1052                let offset = if args.len() > 1 {
1053                    match self.evaluate(&args[1], row_index)? {
1054                        DataValue::Integer(i) => i as i32,
1055                        _ => return Err(anyhow!("LAG offset must be an integer")),
1056                    }
1057                } else {
1058                    1
1059                };
1060
1061                let offset_start = Instant::now();
1062                // Get value at offset
1063                let value = context
1064                    .get_offset_value(row_index, -offset, &column.name)
1065                    .unwrap_or(DataValue::Null);
1066
1067                debug!(
1068                    "LAG offset access took {:.2}μs (offset={})",
1069                    offset_start.elapsed().as_micros(),
1070                    offset
1071                );
1072
1073                Ok(value)
1074            }
1075            "LEAD" => {
1076                // LEAD(column, offset, default)
1077                if args.is_empty() {
1078                    return Err(anyhow!("LEAD requires at least 1 argument"));
1079                }
1080
1081                // Get column name
1082                let column = match &args[0] {
1083                    SqlExpression::Column(col) => col.clone(),
1084                    _ => return Err(anyhow!("LEAD first argument must be a column")),
1085                };
1086
1087                // Get offset (default 1)
1088                let offset = if args.len() > 1 {
1089                    match self.evaluate(&args[1], row_index)? {
1090                        DataValue::Integer(i) => i as i32,
1091                        _ => return Err(anyhow!("LEAD offset must be an integer")),
1092                    }
1093                } else {
1094                    1
1095                };
1096
1097                let offset_start = Instant::now();
1098                // Get value at offset
1099                let value = context
1100                    .get_offset_value(row_index, offset, &column.name)
1101                    .unwrap_or(DataValue::Null);
1102
1103                debug!(
1104                    "LEAD offset access took {:.2}μs (offset={})",
1105                    offset_start.elapsed().as_micros(),
1106                    offset
1107                );
1108
1109                Ok(value)
1110            }
1111            "ROW_NUMBER" => {
1112                // ROW_NUMBER() - no arguments
1113                Ok(DataValue::Integer(context.get_row_number(row_index) as i64))
1114            }
1115            "RANK" => {
1116                // RANK() - no arguments
1117                Ok(DataValue::Integer(context.get_rank(row_index)))
1118            }
1119            "DENSE_RANK" => {
1120                // DENSE_RANK() - no arguments
1121                Ok(DataValue::Integer(context.get_dense_rank(row_index)))
1122            }
1123            "FIRST_VALUE" => {
1124                // FIRST_VALUE(column) OVER (... ROWS ...)
1125                if args.is_empty() {
1126                    return Err(anyhow!("FIRST_VALUE requires 1 argument"));
1127                }
1128
1129                let column = match &args[0] {
1130                    SqlExpression::Column(col) => col.clone(),
1131                    _ => return Err(anyhow!("FIRST_VALUE argument must be a column")),
1132                };
1133
1134                // Use frame-aware version if frame is specified
1135                if context.has_frame() {
1136                    Ok(context
1137                        .get_frame_first_value(row_index, &column.name)
1138                        .unwrap_or(DataValue::Null))
1139                } else {
1140                    Ok(context
1141                        .get_first_value(row_index, &column.name)
1142                        .unwrap_or(DataValue::Null))
1143                }
1144            }
1145            "LAST_VALUE" => {
1146                // LAST_VALUE(column) OVER (... ROWS ...)
1147                if args.is_empty() {
1148                    return Err(anyhow!("LAST_VALUE requires 1 argument"));
1149                }
1150
1151                let column = match &args[0] {
1152                    SqlExpression::Column(col) => col.clone(),
1153                    _ => return Err(anyhow!("LAST_VALUE argument must be a column")),
1154                };
1155
1156                // Use frame-aware version if frame is specified
1157                if context.has_frame() {
1158                    Ok(context
1159                        .get_frame_last_value(row_index, &column.name)
1160                        .unwrap_or(DataValue::Null))
1161                } else {
1162                    Ok(context
1163                        .get_last_value(row_index, &column.name)
1164                        .unwrap_or(DataValue::Null))
1165                }
1166            }
1167            "SUM" => {
1168                // SUM(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1169                if args.is_empty() {
1170                    return Err(anyhow!("SUM requires 1 argument"));
1171                }
1172
1173                let column = match &args[0] {
1174                    SqlExpression::Column(col) => col.clone(),
1175                    _ => return Err(anyhow!("SUM argument must be a column")),
1176                };
1177
1178                // Use frame-aware sum if frame is specified, otherwise use partition sum
1179                if context.has_frame() {
1180                    Ok(context
1181                        .get_frame_sum(row_index, &column.name)
1182                        .unwrap_or(DataValue::Null))
1183                } else {
1184                    Ok(context
1185                        .get_partition_sum(row_index, &column.name)
1186                        .unwrap_or(DataValue::Null))
1187                }
1188            }
1189            "AVG" => {
1190                // AVG(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1191                if args.is_empty() {
1192                    return Err(anyhow!("AVG requires 1 argument"));
1193                }
1194
1195                let column = match &args[0] {
1196                    SqlExpression::Column(col) => col.clone(),
1197                    _ => return Err(anyhow!("AVG argument must be a column")),
1198                };
1199
1200                // Use frame-aware avg if frame is specified, otherwise use partition avg
1201                if context.has_frame() {
1202                    Ok(context
1203                        .get_frame_avg(row_index, &column.name)
1204                        .unwrap_or(DataValue::Null))
1205                } else {
1206                    Ok(context
1207                        .get_partition_avg(row_index, &column.name)
1208                        .unwrap_or(DataValue::Null))
1209                }
1210            }
1211            "STDDEV" | "STDEV" => {
1212                // STDDEV(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1213                if args.is_empty() {
1214                    return Err(anyhow!("STDDEV requires 1 argument"));
1215                }
1216
1217                let column = match &args[0] {
1218                    SqlExpression::Column(col) => col.clone(),
1219                    _ => return Err(anyhow!("STDDEV argument must be a column")),
1220                };
1221
1222                Ok(context
1223                    .get_frame_stddev(row_index, &column.name)
1224                    .unwrap_or(DataValue::Null))
1225            }
1226            "VARIANCE" | "VAR" => {
1227                // VARIANCE(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1228                if args.is_empty() {
1229                    return Err(anyhow!("VARIANCE requires 1 argument"));
1230                }
1231
1232                let column = match &args[0] {
1233                    SqlExpression::Column(col) => col.clone(),
1234                    _ => return Err(anyhow!("VARIANCE argument must be a column")),
1235                };
1236
1237                Ok(context
1238                    .get_frame_variance(row_index, &column.name)
1239                    .unwrap_or(DataValue::Null))
1240            }
1241            "MIN" => {
1242                // MIN(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1243                if args.is_empty() {
1244                    return Err(anyhow!("MIN requires 1 argument"));
1245                }
1246
1247                let column = match &args[0] {
1248                    SqlExpression::Column(col) => col.clone(),
1249                    _ => return Err(anyhow!("MIN argument must be a column")),
1250                };
1251
1252                let frame_rows = context.get_frame_rows(row_index);
1253                if frame_rows.is_empty() {
1254                    return Ok(DataValue::Null);
1255                }
1256
1257                let source_table = context.source();
1258                let col_idx = source_table
1259                    .get_column_index(&column.name)
1260                    .ok_or_else(|| anyhow!("Column '{}' not found", column.name))?;
1261
1262                let mut min_value: Option<DataValue> = None;
1263                for &row_idx in &frame_rows {
1264                    if let Some(value) = source_table.get_value(row_idx, col_idx) {
1265                        if !matches!(value, DataValue::Null) {
1266                            match &min_value {
1267                                None => min_value = Some(value.clone()),
1268                                Some(current_min) => {
1269                                    if value < current_min {
1270                                        min_value = Some(value.clone());
1271                                    }
1272                                }
1273                            }
1274                        }
1275                    }
1276                }
1277
1278                Ok(min_value.unwrap_or(DataValue::Null))
1279            }
1280            "MAX" => {
1281                // MAX(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1282                if args.is_empty() {
1283                    return Err(anyhow!("MAX requires 1 argument"));
1284                }
1285
1286                let column = match &args[0] {
1287                    SqlExpression::Column(col) => col.clone(),
1288                    _ => return Err(anyhow!("MAX argument must be a column")),
1289                };
1290
1291                let frame_rows = context.get_frame_rows(row_index);
1292                if frame_rows.is_empty() {
1293                    return Ok(DataValue::Null);
1294                }
1295
1296                let source_table = context.source();
1297                let col_idx = source_table
1298                    .get_column_index(&column.name)
1299                    .ok_or_else(|| anyhow!("Column '{}' not found", column.name))?;
1300
1301                let mut max_value: Option<DataValue> = None;
1302                for &row_idx in &frame_rows {
1303                    if let Some(value) = source_table.get_value(row_idx, col_idx) {
1304                        if !matches!(value, DataValue::Null) {
1305                            match &max_value {
1306                                None => max_value = Some(value.clone()),
1307                                Some(current_max) => {
1308                                    if value > current_max {
1309                                        max_value = Some(value.clone());
1310                                    }
1311                                }
1312                            }
1313                        }
1314                    }
1315                }
1316
1317                Ok(max_value.unwrap_or(DataValue::Null))
1318            }
1319            "COUNT" => {
1320                // COUNT(*) or COUNT(column) OVER (PARTITION BY ... ROWS n PRECEDING)
1321                // Use frame-aware count if frame is specified, otherwise use partition count
1322
1323                if args.is_empty() {
1324                    // COUNT(*) OVER (...)
1325                    if context.has_frame() {
1326                        Ok(context
1327                            .get_frame_count(row_index, None)
1328                            .unwrap_or(DataValue::Null))
1329                    } else {
1330                        Ok(context
1331                            .get_partition_count(row_index, None)
1332                            .unwrap_or(DataValue::Null))
1333                    }
1334                } else {
1335                    // Check for COUNT(*)
1336                    let column = match &args[0] {
1337                        SqlExpression::Column(col) => {
1338                            if col.name == "*" {
1339                                // COUNT(*) - count all rows
1340                                if context.has_frame() {
1341                                    return Ok(context
1342                                        .get_frame_count(row_index, None)
1343                                        .unwrap_or(DataValue::Null));
1344                                } else {
1345                                    return Ok(context
1346                                        .get_partition_count(row_index, None)
1347                                        .unwrap_or(DataValue::Null));
1348                                }
1349                            }
1350                            col.clone()
1351                        }
1352                        SqlExpression::StringLiteral(s) if s == "*" => {
1353                            // COUNT(*) as StringLiteral
1354                            if context.has_frame() {
1355                                return Ok(context
1356                                    .get_frame_count(row_index, None)
1357                                    .unwrap_or(DataValue::Null));
1358                            } else {
1359                                return Ok(context
1360                                    .get_partition_count(row_index, None)
1361                                    .unwrap_or(DataValue::Null));
1362                            }
1363                        }
1364                        _ => return Err(anyhow!("COUNT argument must be a column or *")),
1365                    };
1366
1367                    // COUNT(column) - count non-null values
1368                    if context.has_frame() {
1369                        Ok(context
1370                            .get_frame_count(row_index, Some(&column.name))
1371                            .unwrap_or(DataValue::Null))
1372                    } else {
1373                        Ok(context
1374                            .get_partition_count(row_index, Some(&column.name))
1375                            .unwrap_or(DataValue::Null))
1376                    }
1377                }
1378            }
1379            _ => Err(anyhow!("Unknown window function: {}", name)),
1380        };
1381
1382        let eval_time = eval_start.elapsed();
1383
1384        info!(
1385            "{} (built-in) evaluation: total={:.2}μs, context={:.2}μs, eval={:.2}μs",
1386            name_upper,
1387            func_start.elapsed().as_micros(),
1388            context_time.as_micros(),
1389            eval_time.as_micros()
1390        );
1391
1392        result
1393    }
1394
1395    /// Evaluate a method call on a column (e.g., `column.Trim()`)
1396    fn evaluate_method_call(
1397        &mut self,
1398        object: &str,
1399        method: &str,
1400        args: &[SqlExpression],
1401        row_index: usize,
1402    ) -> Result<DataValue> {
1403        // Get column value
1404        let col_index = self.table.get_column_index(object).ok_or_else(|| {
1405            let suggestion = self.find_similar_column(object);
1406            match suggestion {
1407                Some(similar) => {
1408                    anyhow!("Column '{}' not found. Did you mean '{}'?", object, similar)
1409                }
1410                None => anyhow!("Column '{}' not found", object),
1411            }
1412        })?;
1413
1414        let cell_value = self.table.get_value(row_index, col_index).cloned();
1415
1416        self.evaluate_method_on_value(
1417            &cell_value.unwrap_or(DataValue::Null),
1418            method,
1419            args,
1420            row_index,
1421        )
1422    }
1423
1424    /// Evaluate a method on a value
1425    fn evaluate_method_on_value(
1426        &mut self,
1427        value: &DataValue,
1428        method: &str,
1429        args: &[SqlExpression],
1430        row_index: usize,
1431    ) -> Result<DataValue> {
1432        // Method-call syntax (`x.Method(...)`) dispatches through the method
1433        // registry first, so a function can give its C#-style method form
1434        // different semantics from its SQL function form (e.g. SUBSTRING is
1435        // 1-based as a function but `.Substring()` is 0-based like .NET).
1436        // The default `evaluate_method` just prepends the receiver and calls
1437        // `evaluate`, so this is behavior-preserving for every other method.
1438        if let Some(method_fn) = self.function_registry.get_method(method) {
1439            let mut method_args = Vec::with_capacity(args.len());
1440            for arg in args {
1441                method_args.push(self.evaluate(arg, row_index)?);
1442            }
1443            return method_fn.evaluate_method(value, &method_args);
1444        }
1445
1446        // Otherwise, proxy the method through the function registry.
1447        // Many string methods have corresponding functions (TRIM, LENGTH, CONTAINS, etc.)
1448
1449        // Map method names to function names (case-insensitive matching)
1450        let function_name = match method.to_lowercase().as_str() {
1451            "trim" => "TRIM",
1452            "trimstart" | "trimbegin" => "TRIMSTART",
1453            "trimend" => "TRIMEND",
1454            "length" | "len" => "LENGTH",
1455            "contains" => "CONTAINS",
1456            "startswith" => "STARTSWITH",
1457            "endswith" => "ENDSWITH",
1458            "indexof" => "INDEXOF",
1459            _ => method, // Try the method name as-is
1460        };
1461
1462        // Check if we have this function in the registry
1463        if self.function_registry.get(function_name).is_some() {
1464            debug!(
1465                "Proxying method '{}' through function registry as '{}'",
1466                method, function_name
1467            );
1468
1469            // Prepare arguments: receiver is the first argument, followed by method args
1470            let mut func_args = vec![value.clone()];
1471
1472            // Evaluate method arguments and add them
1473            for arg in args {
1474                func_args.push(self.evaluate(arg, row_index)?);
1475            }
1476
1477            // Get the function and call it
1478            let func = self.function_registry.get(function_name).unwrap();
1479            return func.evaluate(&func_args);
1480        }
1481
1482        // If not in registry, the method is not supported
1483        // All methods should be registered in the function registry
1484        Err(anyhow!(
1485            "Method '{}' not found. It should be registered in the function registry.",
1486            method
1487        ))
1488    }
1489
1490    /// Evaluate a CASE expression
1491    fn evaluate_case_expression(
1492        &mut self,
1493        when_branches: &[crate::sql::recursive_parser::WhenBranch],
1494        else_branch: &Option<Box<SqlExpression>>,
1495        row_index: usize,
1496    ) -> Result<DataValue> {
1497        debug!(
1498            "ArithmeticEvaluator: evaluating CASE expression for row {}",
1499            row_index
1500        );
1501
1502        // Evaluate each WHEN condition in order
1503        for branch in when_branches {
1504            // Evaluate the condition as a boolean
1505            let condition_result = self.evaluate_condition_as_bool(&branch.condition, row_index)?;
1506
1507            if condition_result {
1508                debug!("CASE: WHEN condition matched, evaluating result expression");
1509                return self.evaluate(&branch.result, row_index);
1510            }
1511        }
1512
1513        // If no WHEN condition matched, evaluate ELSE clause (or return NULL)
1514        if let Some(else_expr) = else_branch {
1515            debug!("CASE: No WHEN matched, evaluating ELSE expression");
1516            self.evaluate(else_expr, row_index)
1517        } else {
1518            debug!("CASE: No WHEN matched and no ELSE, returning NULL");
1519            Ok(DataValue::Null)
1520        }
1521    }
1522
1523    /// Evaluate a simple CASE expression
1524    fn evaluate_simple_case_expression(
1525        &mut self,
1526        expr: &Box<SqlExpression>,
1527        when_branches: &[crate::sql::parser::ast::SimpleWhenBranch],
1528        else_branch: &Option<Box<SqlExpression>>,
1529        row_index: usize,
1530    ) -> Result<DataValue> {
1531        debug!(
1532            "ArithmeticEvaluator: evaluating simple CASE expression for row {}",
1533            row_index
1534        );
1535
1536        // Evaluate the main expression once
1537        let case_value = self.evaluate(expr, row_index)?;
1538        debug!("Simple CASE: evaluated expression to {:?}", case_value);
1539
1540        // Compare against each WHEN value in order
1541        for branch in when_branches {
1542            // Evaluate the WHEN value
1543            let when_value = self.evaluate(&branch.value, row_index)?;
1544
1545            // Check for equality
1546            if self.values_equal(&case_value, &when_value)? {
1547                debug!("Simple CASE: WHEN value matched, evaluating result expression");
1548                return self.evaluate(&branch.result, row_index);
1549            }
1550        }
1551
1552        // If no WHEN value matched, evaluate ELSE clause (or return NULL)
1553        if let Some(else_expr) = else_branch {
1554            debug!("Simple CASE: No WHEN matched, evaluating ELSE expression");
1555            self.evaluate(else_expr, row_index)
1556        } else {
1557            debug!("Simple CASE: No WHEN matched and no ELSE, returning NULL");
1558            Ok(DataValue::Null)
1559        }
1560    }
1561
1562    /// Check if two DataValues are equal
1563    fn values_equal(&self, left: &DataValue, right: &DataValue) -> Result<bool> {
1564        match (left, right) {
1565            (DataValue::Null, DataValue::Null) => Ok(true),
1566            (DataValue::Null, _) | (_, DataValue::Null) => Ok(false),
1567            (DataValue::Integer(a), DataValue::Integer(b)) => Ok(a == b),
1568            (DataValue::Float(a), DataValue::Float(b)) => Ok((a - b).abs() < f64::EPSILON),
1569            (DataValue::String(a), DataValue::String(b)) => Ok(a == b),
1570            (DataValue::Boolean(a), DataValue::Boolean(b)) => Ok(a == b),
1571            (DataValue::DateTime(a), DataValue::DateTime(b)) => Ok(a == b),
1572            // Type coercion for numeric comparisons
1573            (DataValue::Integer(a), DataValue::Float(b)) => {
1574                Ok((*a as f64 - b).abs() < f64::EPSILON)
1575            }
1576            (DataValue::Float(a), DataValue::Integer(b)) => {
1577                Ok((a - *b as f64).abs() < f64::EPSILON)
1578            }
1579            _ => Ok(false),
1580        }
1581    }
1582
1583    /// Helper method to evaluate an expression as a boolean (for CASE WHEN conditions)
1584    fn evaluate_condition_as_bool(
1585        &mut self,
1586        expr: &SqlExpression,
1587        row_index: usize,
1588    ) -> Result<bool> {
1589        let value = self.evaluate(expr, row_index)?;
1590
1591        match value {
1592            DataValue::Boolean(b) => Ok(b),
1593            DataValue::Integer(i) => Ok(i != 0),
1594            DataValue::Float(f) => Ok(f != 0.0),
1595            DataValue::Null => Ok(false),
1596            DataValue::String(s) => Ok(!s.is_empty()),
1597            DataValue::InternedString(s) => Ok(!s.is_empty()),
1598            _ => Ok(true), // Other types are considered truthy
1599        }
1600    }
1601
1602    /// Evaluate a DATETIME constructor expression
1603    fn evaluate_datetime_constructor(
1604        &self,
1605        year: i32,
1606        month: u32,
1607        day: u32,
1608        hour: Option<u32>,
1609        minute: Option<u32>,
1610        second: Option<u32>,
1611    ) -> Result<DataValue> {
1612        use chrono::{NaiveDate, TimeZone, Utc};
1613
1614        // Create a NaiveDate
1615        let date = NaiveDate::from_ymd_opt(year, month, day)
1616            .ok_or_else(|| anyhow!("Invalid date: {}-{}-{}", year, month, day))?;
1617
1618        // Create datetime with provided time components or defaults
1619        let hour = hour.unwrap_or(0);
1620        let minute = minute.unwrap_or(0);
1621        let second = second.unwrap_or(0);
1622
1623        let naive_datetime = date
1624            .and_hms_opt(hour, minute, second)
1625            .ok_or_else(|| anyhow!("Invalid time: {}:{}:{}", hour, minute, second))?;
1626
1627        // Convert to UTC DateTime
1628        let datetime = Utc.from_utc_datetime(&naive_datetime);
1629
1630        // Format as string with milliseconds
1631        let datetime_str = datetime.format("%Y-%m-%d %H:%M:%S%.3f").to_string();
1632        Ok(DataValue::String(datetime_str))
1633    }
1634
1635    /// Evaluate a DATETIME.TODAY constructor expression
1636    fn evaluate_datetime_today(
1637        &self,
1638        hour: Option<u32>,
1639        minute: Option<u32>,
1640        second: Option<u32>,
1641    ) -> Result<DataValue> {
1642        use chrono::{TimeZone, Utc};
1643
1644        // Get today's date in UTC
1645        let today = Utc::now().date_naive();
1646
1647        // Create datetime with provided time components or defaults
1648        let hour = hour.unwrap_or(0);
1649        let minute = minute.unwrap_or(0);
1650        let second = second.unwrap_or(0);
1651
1652        let naive_datetime = today
1653            .and_hms_opt(hour, minute, second)
1654            .ok_or_else(|| anyhow!("Invalid time: {}:{}:{}", hour, minute, second))?;
1655
1656        // Convert to UTC DateTime
1657        let datetime = Utc.from_utc_datetime(&naive_datetime);
1658
1659        // Format as string with milliseconds
1660        let datetime_str = datetime.format("%Y-%m-%d %H:%M:%S%.3f").to_string();
1661        Ok(DataValue::String(datetime_str))
1662    }
1663}
1664
1665#[cfg(test)]
1666mod tests {
1667    use super::*;
1668    use crate::data::datatable::{DataColumn, DataRow};
1669
1670    fn create_test_table() -> DataTable {
1671        let mut table = DataTable::new("test");
1672        table.add_column(DataColumn::new("a"));
1673        table.add_column(DataColumn::new("b"));
1674        table.add_column(DataColumn::new("c"));
1675
1676        table
1677            .add_row(DataRow::new(vec![
1678                DataValue::Integer(10),
1679                DataValue::Float(2.5),
1680                DataValue::Integer(4),
1681            ]))
1682            .unwrap();
1683
1684        table
1685    }
1686
1687    #[test]
1688    fn test_evaluate_column() {
1689        let table = create_test_table();
1690        let mut evaluator = ArithmeticEvaluator::new(&table);
1691
1692        let expr = SqlExpression::Column(ColumnRef::unquoted("a".to_string()));
1693        let result = evaluator.evaluate(&expr, 0).unwrap();
1694        assert_eq!(result, DataValue::Integer(10));
1695    }
1696
1697    #[test]
1698    fn test_evaluate_between_column_in_range() {
1699        let table = create_test_table();
1700        let mut evaluator = ArithmeticEvaluator::new(&table);
1701
1702        // column 'a' is 10 — 5 <= 10 <= 20 is true
1703        let expr = SqlExpression::Between {
1704            expr: Box::new(SqlExpression::Column(ColumnRef::unquoted("a".to_string()))),
1705            lower: Box::new(SqlExpression::NumberLiteral("5".to_string())),
1706            upper: Box::new(SqlExpression::NumberLiteral("20".to_string())),
1707        };
1708        assert_eq!(
1709            evaluator.evaluate(&expr, 0).unwrap(),
1710            DataValue::Boolean(true)
1711        );
1712    }
1713
1714    #[test]
1715    fn test_evaluate_between_column_out_of_range() {
1716        let table = create_test_table();
1717        let mut evaluator = ArithmeticEvaluator::new(&table);
1718
1719        // column 'a' is 10 — 11 <= 10 <= 20 is false
1720        let expr = SqlExpression::Between {
1721            expr: Box::new(SqlExpression::Column(ColumnRef::unquoted("a".to_string()))),
1722            lower: Box::new(SqlExpression::NumberLiteral("11".to_string())),
1723            upper: Box::new(SqlExpression::NumberLiteral("20".to_string())),
1724        };
1725        assert_eq!(
1726            evaluator.evaluate(&expr, 0).unwrap(),
1727            DataValue::Boolean(false)
1728        );
1729    }
1730
1731    #[test]
1732    fn test_evaluate_between_endpoints_inclusive() {
1733        let table = create_test_table();
1734        let mut evaluator = ArithmeticEvaluator::new(&table);
1735
1736        // column 'a' is 10 — 10 <= 10 <= 10 is true (both endpoints inclusive)
1737        let expr = SqlExpression::Between {
1738            expr: Box::new(SqlExpression::Column(ColumnRef::unquoted("a".to_string()))),
1739            lower: Box::new(SqlExpression::NumberLiteral("10".to_string())),
1740            upper: Box::new(SqlExpression::NumberLiteral("10".to_string())),
1741        };
1742        assert_eq!(
1743            evaluator.evaluate(&expr, 0).unwrap(),
1744            DataValue::Boolean(true)
1745        );
1746    }
1747
1748    #[test]
1749    fn test_evaluate_number_literal() {
1750        let table = create_test_table();
1751        let mut evaluator = ArithmeticEvaluator::new(&table);
1752
1753        let expr = SqlExpression::NumberLiteral("42".to_string());
1754        let result = evaluator.evaluate(&expr, 0).unwrap();
1755        assert_eq!(result, DataValue::Integer(42));
1756
1757        let expr = SqlExpression::NumberLiteral("3.14".to_string());
1758        let result = evaluator.evaluate(&expr, 0).unwrap();
1759        assert_eq!(result, DataValue::Float(3.14));
1760    }
1761
1762    #[test]
1763    fn test_add_values() {
1764        let table = create_test_table();
1765        let mut evaluator = ArithmeticEvaluator::new(&table);
1766
1767        // Integer + Integer
1768        let result = evaluator
1769            .add_values(&DataValue::Integer(5), &DataValue::Integer(3))
1770            .unwrap();
1771        assert_eq!(result, DataValue::Integer(8));
1772
1773        // Integer + Float
1774        let result = evaluator
1775            .add_values(&DataValue::Integer(5), &DataValue::Float(2.5))
1776            .unwrap();
1777        assert_eq!(result, DataValue::Float(7.5));
1778    }
1779
1780    #[test]
1781    fn test_multiply_values() {
1782        let table = create_test_table();
1783        let mut evaluator = ArithmeticEvaluator::new(&table);
1784
1785        // Integer * Float
1786        let result = evaluator
1787            .multiply_values(&DataValue::Integer(4), &DataValue::Float(2.5))
1788            .unwrap();
1789        assert_eq!(result, DataValue::Float(10.0));
1790    }
1791
1792    #[test]
1793    fn test_divide_values() {
1794        let table = create_test_table();
1795        let mut evaluator = ArithmeticEvaluator::new(&table);
1796
1797        // Exact division
1798        let result = evaluator
1799            .divide_values(&DataValue::Integer(10), &DataValue::Integer(2))
1800            .unwrap();
1801        assert_eq!(result, DataValue::Integer(5));
1802
1803        // Non-exact division
1804        let result = evaluator
1805            .divide_values(&DataValue::Integer(10), &DataValue::Integer(3))
1806            .unwrap();
1807        assert_eq!(result, DataValue::Float(10.0 / 3.0));
1808    }
1809
1810    #[test]
1811    fn test_division_by_zero() {
1812        let table = create_test_table();
1813        let mut evaluator = ArithmeticEvaluator::new(&table);
1814
1815        let result = evaluator.divide_values(&DataValue::Integer(10), &DataValue::Integer(0));
1816        assert!(result.is_err());
1817        assert!(result.unwrap_err().to_string().contains("Division by zero"));
1818    }
1819
1820    #[test]
1821    fn test_binary_op_expression() {
1822        let table = create_test_table();
1823        let mut evaluator = ArithmeticEvaluator::new(&table);
1824
1825        // a * b where a=10, b=2.5
1826        let expr = SqlExpression::BinaryOp {
1827            left: Box::new(SqlExpression::Column(ColumnRef::unquoted("a".to_string()))),
1828            op: "*".to_string(),
1829            right: Box::new(SqlExpression::Column(ColumnRef::unquoted("b".to_string()))),
1830        };
1831
1832        let result = evaluator.evaluate(&expr, 0).unwrap();
1833        assert_eq!(result, DataValue::Float(25.0));
1834    }
1835}
sql_cli/data/arithmetic_evaluator.rs

sql_cli/data/
arithmetic_evaluator.rs
