ddb_core/engine/
executor.rs

1// Query executor for SELECT statements
2use crate::config::Table;
3use crate::engine::{evaluator::Evaluator, row::Row};
4use crate::error::{DdbError, Result};
5use crate::file_io::CsvReader;
6use crate::functions::Value;
7use crate::parser::{OrderDirection, SelectColumn, SelectStatement};
8
9pub struct QueryExecutor {
10    evaluator: Evaluator,
11}
12
13impl QueryExecutor {
14    pub fn new() -> Self {
15        QueryExecutor {
16            evaluator: Evaluator::new(),
17        }
18    }
19
20    /// Execute a SELECT statement
21    pub fn execute_select(&self, stmt: &SelectStatement, table: &Table) -> Result<Vec<Row>> {
22        // Check if query uses aggregate functions
23        let has_aggregates = self.has_aggregate_functions(&stmt.columns);
24
25        if has_aggregates {
26            return self.execute_aggregate_select(stmt, table);
27        }
28
29        // Read data from file
30        let mut reader = CsvReader::new(&table.data_file, table.field_delimiter, true)?;
31
32        let mut rows = Vec::new();
33
34        // Filter rows based on WHERE clause
35        while let Some(row) = reader.next_row()? {
36            if let Some(where_expr) = &stmt.where_clause {
37                let result = self.evaluator.evaluate(where_expr, &row)?;
38                let matches = self.to_boolean(&result)?;
39                if !matches {
40                    continue; // Skip this row
41                }
42            }
43
44            // Project columns (SELECT clause)
45            let projected_row = self.project_row(&row, &stmt.columns)?;
46            rows.push(projected_row);
47        }
48
49        // Apply DISTINCT
50        if stmt.distinct {
51            rows = self.deduplicate_rows(rows);
52        }
53
54        // Apply ORDER BY
55        if !stmt.order_by.is_empty() {
56            rows = self.sort_rows(rows, &stmt.order_by)?;
57        }
58
59        // Apply LIMIT
60        if let Some(limit) = &stmt.limit {
61            let start = limit.offset.unwrap_or(0);
62            rows = rows.into_iter().skip(start).take(limit.count).collect();
63        }
64
65        Ok(rows)
66    }
67
68    /// Check if SELECT columns contain aggregate functions
69    fn has_aggregate_functions(&self, columns: &[SelectColumn]) -> bool {
70        columns.iter().any(|col| match col {
71            SelectColumn::Function { name, .. } => {
72                matches!(
73                    name.to_uppercase().as_str(),
74                    "COUNT" | "SUM" | "AVG" | "MIN" | "MAX" | "STDDEV" | "VARIANCE" | "GROUP_CONCAT"
75                )
76            }
77            _ => false,
78        })
79    }
80
81    /// Execute SELECT with aggregate functions
82    fn execute_aggregate_select(&self, stmt: &SelectStatement, table: &Table) -> Result<Vec<Row>> {
83        // Read and filter all rows first
84        let mut reader = CsvReader::new(&table.data_file, table.field_delimiter, true)?;
85        let mut all_rows = Vec::new();
86
87        while let Some(row) = reader.next_row()? {
88            if let Some(where_expr) = &stmt.where_clause {
89                let result = self.evaluator.evaluate(where_expr, &row)?;
90                let matches = self.to_boolean(&result)?;
91                if !matches {
92                    continue;
93                }
94            }
95            all_rows.push(row);
96        }
97
98        // Build result row with aggregates
99        let mut result_row = Row::new();
100
101        for col_spec in &stmt.columns {
102            match col_spec {
103                SelectColumn::Function { name, args, alias } => {
104                    // Special handling for COUNT(*)
105                    let aggregate_result = if self.is_count_star(name, args) {
106                        // COUNT(*) - return row count directly
107                        Value::Integer(all_rows.len() as i64)
108                    } else {
109                        // Collect values for this column
110                        let values: Result<Vec<Value>> = all_rows
111                            .iter()
112                            .map(|row| {
113                                if args.len() == 1 {
114                                    self.evaluator.evaluate(&args[0], row)
115                                } else {
116                                    Err(DdbError::ExecutionError(
117                                        "Aggregate functions require exactly one argument".to_string(),
118                                    ))
119                                }
120                            })
121                            .collect();
122
123                        let values = values?;
124
125                        // Apply aggregate function
126                        match name.to_uppercase().as_str() {
127                            "COUNT" => crate::functions::aggregate::count(&values, true)?,
128                            "SUM" => crate::functions::aggregate::sum(&values)?,
129                            "AVG" => crate::functions::aggregate::avg(&values)?,
130                            "MIN" => crate::functions::aggregate::min(&values)?,
131                            "MAX" => crate::functions::aggregate::max(&values)?,
132                            "STDDEV" | "STDDEV_POP" => crate::functions::aggregate::stddev_pop(&values)?,
133                            "VARIANCE" | "VAR_POP" => crate::functions::aggregate::var_pop(&values)?,
134                            _ => {
135                                return Err(DdbError::ExecutionError(format!(
136                                    "Unknown aggregate function: {}",
137                                    name
138                                )))
139                            }
140                        }
141                    };
142
143                    let default_name = format!("{}(...)", name);
144                    let output_name = alias.as_ref().unwrap_or(&default_name);
145                    result_row.set(output_name, aggregate_result);
146                }
147                SelectColumn::Column { name, alias } => {
148                    // Non-aggregate column in aggregate query - use first value
149                    // (This is technically not standard SQL without GROUP BY, but we'll allow it)
150                    let value = all_rows
151                        .first()
152                        .and_then(|row| row.get(name).cloned())
153                        .unwrap_or(Value::Null);
154                    let output_name = alias.as_ref().unwrap_or(name);
155                    result_row.set(output_name, value);
156                }
157                SelectColumn::Wildcard => {
158                    return Err(DdbError::ExecutionError(
159                        "Cannot use * with aggregate functions (use COUNT(*) instead)".to_string(),
160                    ));
161                }
162            }
163        }
164
165        Ok(vec![result_row])
166    }
167
168    /// Check if this is COUNT(*)
169    fn is_count_star(&self, name: &str, args: &[crate::parser::Expression]) -> bool {
170        name.to_uppercase() == "COUNT"
171            && args.len() == 1
172            && matches!(&args[0], crate::parser::Expression::Column(col) if col == "*")
173    }
174
175    /// Project columns from a row based on SELECT clause
176    fn project_row(&self, row: &Row, columns: &[SelectColumn]) -> Result<Row> {
177        let mut new_row = Row::new();
178
179        for col_spec in columns {
180            match col_spec {
181                SelectColumn::Wildcard => {
182                    // Include all columns
183                    for (name, value) in row.as_map() {
184                        new_row.set(name, value.clone());
185                    }
186                }
187                SelectColumn::Column { name, alias } => {
188                    // Check if this is a system variable (@@VARIABLE)
189                    let value = if name.starts_with("@@") {
190                        let var_name = &name[2..];
191                        let var_expr = crate::parser::Expression::SystemVariable(var_name.to_string());
192                        self.evaluator.evaluate(&var_expr, row)?
193                    } else {
194                        row.get(name)
195                            .cloned()
196                            .ok_or_else(|| DdbError::ColumnNotFound(name.clone()))?
197                    };
198                    let output_name = alias.as_ref().unwrap_or(name);
199                    new_row.set(output_name, value);
200                }
201                SelectColumn::Function { name, args, alias } => {
202                    // Evaluate function
203                    let func_expr = crate::parser::Expression::Function {
204                        name: name.clone(),
205                        args: args.clone(),
206                    };
207                    let result = self.evaluator.evaluate(&func_expr, row)?;
208
209                    let default_name = format!("{}(...)", name);
210                    let output_name = alias.as_ref().unwrap_or(&default_name);
211                    new_row.set(output_name, result);
212                }
213            }
214        }
215
216        Ok(new_row)
217    }
218
219    /// Remove duplicate rows
220    fn deduplicate_rows(&self, mut rows: Vec<Row>) -> Vec<Row> {
221        let mut seen = std::collections::HashSet::new();
222        rows.retain(|row| {
223            let key = self.row_to_key(row);
224            seen.insert(key)
225        });
226        rows
227    }
228
229    /// Convert row to a hashable key for deduplication
230    fn row_to_key(&self, row: &Row) -> Vec<String> {
231        let mut keys: Vec<_> = row.column_names();
232        keys.sort();
233        keys.iter()
234            .map(|name| {
235                row.get(name)
236                    .map(|v| v.to_string())
237                    .unwrap_or_else(|| "NULL".to_string())
238            })
239            .collect()
240    }
241
242    /// Sort rows based on ORDER BY clause
243    fn sort_rows(
244        &self,
245        mut rows: Vec<Row>,
246        order_by: &[crate::parser::OrderByColumn],
247    ) -> Result<Vec<Row>> {
248        rows.sort_by(|a, b| {
249            for order_col in order_by {
250                let a_val = a.get(&order_col.column).cloned().unwrap_or(Value::Null);
251                let b_val = b.get(&order_col.column).cloned().unwrap_or(Value::Null);
252
253                let cmp = self.compare_values(&a_val, &b_val);
254
255                let cmp = match order_col.direction {
256                    OrderDirection::Asc => cmp,
257                    OrderDirection::Desc => cmp.reverse(),
258                };
259
260                if cmp != std::cmp::Ordering::Equal {
261                    return cmp;
262                }
263            }
264            std::cmp::Ordering::Equal
265        });
266
267        Ok(rows)
268    }
269
270    /// Compare two values for sorting
271    fn compare_values(&self, left: &Value, right: &Value) -> std::cmp::Ordering {
272        use std::cmp::Ordering;
273
274        match (left, right) {
275            (Value::Null, Value::Null) => Ordering::Equal,
276            (Value::Null, _) => Ordering::Less,
277            (_, Value::Null) => Ordering::Greater,
278            (Value::Integer(a), Value::Integer(b)) => a.cmp(b),
279            (Value::Float(a), Value::Float(b)) => a.partial_cmp(b).unwrap_or(Ordering::Equal),
280            (Value::Integer(a), Value::Float(b)) => {
281                (*a as f64).partial_cmp(b).unwrap_or(Ordering::Equal)
282            }
283            (Value::Float(a), Value::Integer(b)) => {
284                a.partial_cmp(&(*b as f64)).unwrap_or(Ordering::Equal)
285            }
286            (Value::String(a), Value::String(b)) => a.cmp(b),
287            (Value::Boolean(a), Value::Boolean(b)) => a.cmp(b),
288            (Value::Date(a), Value::Date(b)) => a.cmp(b),
289            (Value::DateTime(a), Value::DateTime(b)) => a.cmp(b),
290            (Value::Time(a), Value::Time(b)) => a.cmp(b),
291            _ => Ordering::Equal,
292        }
293    }
294
295    /// Convert value to boolean
296    fn to_boolean(&self, value: &Value) -> Result<bool> {
297        match value {
298            Value::Boolean(b) => Ok(*b),
299            Value::Integer(i) => Ok(*i != 0),
300            Value::Float(f) => Ok(*f != 0.0),
301            Value::String(s) => Ok(!s.is_empty()),
302            Value::Null => Ok(false),
303            _ => Ok(true),
304        }
305    }
306}
307
308impl Default for QueryExecutor {
309    fn default() -> Self {
310        Self::new()
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317    use crate::parser::{BinaryOperator, Expression, Literal};
318    use std::io::Write;
319    use tempfile::NamedTempFile;
320
321    fn create_test_table() -> (Table, NamedTempFile) {
322        let mut file = NamedTempFile::new().unwrap();
323        writeln!(file, "id,name,age").unwrap();
324        writeln!(file, "1,Alice,30").unwrap();
325        writeln!(file, "2,Bob,25").unwrap();
326        writeln!(file, "3,Charlie,35").unwrap();
327        file.flush().unwrap();
328
329        let table = Table {
330            name: "users".to_string(),
331            database: "test".to_string(),
332            data_file: file.path().to_string_lossy().to_string(),
333            columns: vec![],
334            field_delimiter: ',',
335            data_starts_on: 0,
336            comment_char: None,
337        };
338
339        (table, file)
340    }
341
342    #[test]
343    fn test_execute_select_all() {
344        let (table, _file) = create_test_table();
345        let executor = QueryExecutor::new();
346
347        let stmt = SelectStatement {
348            distinct: false,
349            columns: vec![SelectColumn::Wildcard],
350            from: Some("users".to_string()),
351            where_clause: None,
352            order_by: vec![],
353            limit: None,
354        };
355
356        let results = executor.execute_select(&stmt, &table).unwrap();
357        assert_eq!(results.len(), 3);
358    }
359
360    #[test]
361    fn test_execute_select_with_where() {
362        let (table, _file) = create_test_table();
363        let executor = QueryExecutor::new();
364
365        // WHERE age > 25
366        let where_expr = Expression::BinaryOp {
367            left: Box::new(Expression::Column("age".to_string())),
368            op: BinaryOperator::GreaterThan,
369            right: Box::new(Expression::Literal(Literal::Integer(25))),
370        };
371
372        let stmt = SelectStatement {
373            distinct: false,
374            columns: vec![SelectColumn::Wildcard],
375            from: Some("users".to_string()),
376            where_clause: Some(where_expr),
377            order_by: vec![],
378            limit: None,
379        };
380
381        let results = executor.execute_select(&stmt, &table).unwrap();
382        assert_eq!(results.len(), 2); // Alice (30) and Charlie (35)
383    }
384
385    #[test]
386    fn test_execute_select_with_columns() {
387        let (table, _file) = create_test_table();
388        let executor = QueryExecutor::new();
389
390        let stmt = SelectStatement {
391            distinct: false,
392            columns: vec![
393                SelectColumn::Column {
394                    name: "name".to_string(),
395                    alias: None,
396                },
397                SelectColumn::Column {
398                    name: "age".to_string(),
399                    alias: None,
400                },
401            ],
402            from: Some("users".to_string()),
403            where_clause: None,
404            order_by: vec![],
405            limit: None,
406        };
407
408        let results = executor.execute_select(&stmt, &table).unwrap();
409        assert_eq!(results.len(), 3);
410
411        // Check that only selected columns are present
412        let row = &results[0];
413        assert!(row.get("name").is_some());
414        assert!(row.get("age").is_some());
415        // id should not be present
416        assert!(row.get("id").is_none());
417    }
418
419    #[test]
420    fn test_execute_select_with_limit() {
421        let (table, _file) = create_test_table();
422        let executor = QueryExecutor::new();
423
424        let stmt = SelectStatement {
425            distinct: false,
426            columns: vec![SelectColumn::Wildcard],
427            from: Some("users".to_string()),
428            where_clause: None,
429            order_by: vec![],
430            limit: Some(crate::parser::Limit {
431                offset: None,
432                count: 2,
433            }),
434        };
435
436        let results = executor.execute_select(&stmt, &table).unwrap();
437        assert_eq!(results.len(), 2);
438    }
439
440    #[test]
441    fn test_execute_select_with_order_by() {
442        let (table, _file) = create_test_table();
443        let executor = QueryExecutor::new();
444
445        let stmt = SelectStatement {
446            distinct: false,
447            columns: vec![SelectColumn::Wildcard],
448            from: Some("users".to_string()),
449            where_clause: None,
450            order_by: vec![crate::parser::OrderByColumn {
451                column: "age".to_string(),
452                direction: OrderDirection::Desc,
453            }],
454            limit: None,
455        };
456
457        let results = executor.execute_select(&stmt, &table).unwrap();
458        assert_eq!(results.len(), 3);
459
460        // Check order: Charlie (35), Alice (30), Bob (25)
461        assert_eq!(
462            results[0].get("name"),
463            Some(&Value::String("Charlie".to_string()))
464        );
465        assert_eq!(
466            results[2].get("name"),
467            Some(&Value::String("Bob".to_string()))
468        );
469    }
470}