vibesql_executor/select/executor/
arena_execution.rs

1//! Arena-based SELECT execution for zero-allocation prepared statement execution.
2//!
3//! This module provides arena-based execution of SELECT statements, enabling
4//! zero-allocation query execution for prepared statements with inline
5//! placeholder resolution.
6//!
7//! # Usage
8//!
9//! ```text
10//! use bumpalo::Bump;
11//! use vibesql_parser::arena_parser::ArenaParser;
12//!
13//! let arena = Bump::new();
14//! let stmt = ArenaParser::parse_sql("SELECT * FROM users WHERE id = ?", &arena)?;
15//! let params = &[SqlValue::Integer(42)];
16//! let result = executor.execute_select_arena(stmt, params)?;
17//! ```
18//!
19//! # Performance
20//!
21//! Arena-based execution provides:
22//! - Zero malloc/free overhead per query execution
23//! - Inline placeholder resolution (no AST cloning)
24//! - Direct evaluation without intermediate allocations
25//!
26//! This is particularly beneficial for OLTP workloads with high query rates.
27
28use std::cmp::Ordering;
29use std::collections::HashMap;
30
31use vibesql_ast::arena::{ArenaInterner, Expression as ArenaExpression, ExtendedExpr as ArenaExtendedExpr, SelectItem as ArenaSelectItem, SelectStmt as ArenaSelectStmt};
32use vibesql_storage::Row;
33use vibesql_types::SqlValue;
34
35use super::builder::SelectExecutor;
36use crate::errors::ExecutorError;
37use crate::evaluator::ArenaExpressionEvaluator;
38use crate::evaluator::window::compare_values;
39use crate::schema::CombinedSchema;
40
41impl SelectExecutor<'_> {
42    /// Execute an arena-allocated SELECT statement with inline placeholder resolution.
43    ///
44    /// This method provides zero-allocation query execution for prepared statements.
45    /// Parameters are resolved inline during evaluation, avoiding AST cloning.
46    ///
47    /// # Arguments
48    ///
49    /// * `stmt` - Arena-allocated SELECT statement
50    /// * `params` - Parameter values for placeholder resolution
51    ///
52    /// # Returns
53    ///
54    /// Vector of result rows.
55    ///
56    /// # Limitations
57    ///
58    /// Currently supports simple SELECT queries without:
59    /// - JOINs (single table only)
60    /// - Subqueries
61    /// - Aggregates, GROUP BY, HAVING
62    /// - Window functions
63    /// - CTEs (WITH clause)
64    ///
65    /// For unsupported features, returns an error. Fall back to standard execution.
66    pub fn execute_select_arena<'arena>(
67        &self,
68        stmt: &ArenaSelectStmt<'arena>,
69        params: &[SqlValue],
70        interner: &'arena ArenaInterner<'arena>,
71    ) -> Result<Vec<Row>, ExecutorError> {
72        // Check for unsupported features
73        if stmt.with_clause.is_some() {
74            return Err(ExecutorError::UnsupportedExpression(
75                "Arena execution does not support WITH clause".to_string(),
76            ));
77        }
78
79        if stmt.set_operation.is_some() {
80            return Err(ExecutorError::UnsupportedExpression(
81                "Arena execution does not support set operations".to_string(),
82            ));
83        }
84
85        if stmt.group_by.is_some() || stmt.having.is_some() {
86            return Err(ExecutorError::UnsupportedExpression(
87                "Arena execution does not support GROUP BY/HAVING".to_string(),
88            ));
89        }
90
91        if stmt.distinct {
92            return Err(ExecutorError::UnsupportedExpression(
93                "Arena execution does not support DISTINCT".to_string(),
94            ));
95        }
96
97        // Check for aggregates in select list
98        if self.has_arena_aggregates(&stmt.select_list) {
99            return Err(ExecutorError::UnsupportedExpression(
100                "Arena execution does not support aggregate functions".to_string(),
101            ));
102        }
103
104        // Execute based on FROM clause
105        match &stmt.from {
106            Some(from) => self.execute_arena_with_from(stmt, from, params, interner),
107            None => self.execute_arena_without_from(stmt, params, interner),
108        }
109    }
110
111    /// Execute SELECT without FROM clause (expression-only).
112    fn execute_arena_without_from<'arena>(
113        &self,
114        stmt: &ArenaSelectStmt<'arena>,
115        params: &[SqlValue],
116        interner: &'arena ArenaInterner<'arena>,
117    ) -> Result<Vec<Row>, ExecutorError> {
118        // Create an empty schema and row for expression evaluation
119        let schema = CombinedSchema {
120            table_schemas: HashMap::new(),
121            total_columns: 0,
122        };
123        let empty_row = Row::new(vec![]);
124        let evaluator = ArenaExpressionEvaluator::new(&schema, params, interner);
125
126        // Evaluate each select item
127        let mut values = Vec::with_capacity(stmt.select_list.len());
128        for item in stmt.select_list.iter() {
129            match item {
130                ArenaSelectItem::Expression { expr, .. } => {
131                    let value = evaluator.eval(expr, &empty_row)?;
132                    values.push(value);
133                }
134                ArenaSelectItem::Wildcard { .. } | ArenaSelectItem::QualifiedWildcard { .. } => {
135                    // Wildcard without FROM is typically an error, but we'll skip it
136                    continue;
137                }
138            }
139        }
140
141        // Apply LIMIT/OFFSET (for expression-only, limit defaults to 1 row)
142        let rows = vec![Row::new(values)];
143        Ok(self.apply_arena_limit_offset(rows, stmt.limit, stmt.offset))
144    }
145
146    /// Execute SELECT with FROM clause.
147    fn execute_arena_with_from<'arena>(
148        &self,
149        stmt: &ArenaSelectStmt<'arena>,
150        from: &vibesql_ast::arena::FromClause<'arena>,
151        params: &[SqlValue],
152        interner: &'arena ArenaInterner<'arena>,
153    ) -> Result<Vec<Row>, ExecutorError> {
154        use vibesql_ast::arena::FromClause;
155
156        // Currently only support simple table reference
157        let (table_name, alias) = match from {
158            FromClause::Table { name, alias, .. } => (*name, *alias),
159            FromClause::Join { .. } => {
160                return Err(ExecutorError::UnsupportedExpression(
161                    "Arena execution does not support JOINs yet".to_string(),
162                ));
163            }
164            FromClause::Subquery { .. } => {
165                return Err(ExecutorError::UnsupportedExpression(
166                    "Arena execution does not support subqueries in FROM".to_string(),
167                ));
168            }
169        };
170
171        // Resolve table name symbol to string
172        let table_name_str = interner.resolve(table_name);
173
174        // Get the table
175        let table = self.database.get_table(table_name_str).ok_or_else(|| {
176            ExecutorError::TableNotFound(table_name_str.to_string())
177        })?;
178
179        // Build schema - use alias if provided, otherwise table name
180        let schema_alias_str = alias.map(|a| interner.resolve(a)).unwrap_or(table_name_str);
181        let schema = CombinedSchema::from_table(schema_alias_str.to_string(), table.schema.clone());
182
183        // Create evaluator
184        let evaluator = ArenaExpressionEvaluator::with_database(&schema, params, self.database, interner);
185
186        // Scan table and filter
187        let mut results = Vec::new();
188        for row in table.scan() {
189            // Apply WHERE clause filter
190            if let Some(where_clause) = &stmt.where_clause {
191                let filter_result = evaluator.eval(where_clause, row)?;
192                match filter_result {
193                    SqlValue::Boolean(true) => {}
194                    SqlValue::Boolean(false) | SqlValue::Null => continue,
195                    _ => {
196                        return Err(ExecutorError::TypeError(format!(
197                            "WHERE clause must evaluate to boolean, got {:?}",
198                            filter_result
199                        )));
200                    }
201                }
202            }
203
204            // Project columns
205            let projected = self.project_arena_row(&stmt.select_list, row, &schema, &evaluator, interner)?;
206            results.push(projected);
207
208            // Check timeout periodically
209            if results.len() % 1000 == 0 {
210                self.check_timeout()?;
211            }
212        }
213
214        // Apply ORDER BY if present
215        if let Some(order_by) = &stmt.order_by {
216            self.sort_arena_results(&mut results, order_by.as_slice(), &schema, params, interner)?;
217        }
218
219        // Apply LIMIT/OFFSET
220        Ok(self.apply_arena_limit_offset(results, stmt.limit, stmt.offset))
221    }
222
223    /// Project a row according to the SELECT list.
224    fn project_arena_row<'arena>(
225        &self,
226        select_list: &[ArenaSelectItem<'arena>],
227        row: &Row,
228        schema: &CombinedSchema,
229        evaluator: &ArenaExpressionEvaluator<'_, 'arena>,
230        interner: &'arena ArenaInterner<'arena>,
231    ) -> Result<Row, ExecutorError> {
232        let mut values = Vec::with_capacity(select_list.len());
233
234        for item in select_list.iter() {
235            match item {
236                ArenaSelectItem::Expression { expr, .. } => {
237                    let value = evaluator.eval(expr, row)?;
238                    values.push(value);
239                }
240                ArenaSelectItem::Wildcard { .. } => {
241                    // Unqualified wildcard (*) - expand all columns
242                    values.extend(row.values.iter().cloned());
243                }
244                ArenaSelectItem::QualifiedWildcard { qualifier, .. } => {
245                    // Qualified wildcard (table.*) - expand columns from specific table
246                    let qualifier_str = interner.resolve(*qualifier);
247                    if let Some(&(start, ref tbl_schema)) = schema.table_schemas.get(&qualifier_str.to_lowercase()) {
248                        for i in 0..tbl_schema.columns.len() {
249                            if let Some(val) = row.get(start + i) {
250                                values.push(val.clone());
251                            }
252                        }
253                    } else {
254                        // Table not found - expand all as fallback
255                        values.extend(row.values.iter().cloned());
256                    }
257                }
258            }
259        }
260
261        Ok(Row::new(values))
262    }
263
264    /// Sort results according to ORDER BY clause.
265    fn sort_arena_results<'arena>(
266        &self,
267        results: &mut Vec<Row>,
268        order_by: &[vibesql_ast::arena::OrderByItem<'arena>],
269        schema: &CombinedSchema,
270        params: &[SqlValue],
271        interner: &'arena ArenaInterner<'arena>,
272    ) -> Result<(), ExecutorError> {
273        use vibesql_ast::arena::OrderDirection;
274
275        // Create evaluator for order by expressions
276        let evaluator = ArenaExpressionEvaluator::with_database(schema, params, self.database, interner);
277
278        // Pre-compute order by values for each row to avoid repeated evaluation
279        let mut keyed_rows: Vec<(Vec<SqlValue>, Row)> = results
280            .drain(..)
281            .map(|row| {
282                let keys: Result<Vec<_>, _> = order_by
283                    .iter()
284                    .map(|item| evaluator.eval(&item.expr, &row))
285                    .collect();
286                keys.map(|k| (k, row))
287            })
288            .collect::<Result<_, _>>()?;
289
290        // Sort by keys
291        keyed_rows.sort_by(|(keys_a, _), (keys_b, _)| {
292            for (i, (key_a, key_b)) in keys_a.iter().zip(keys_b.iter()).enumerate() {
293                let cmp = compare_values(key_a, key_b);
294                if cmp != Ordering::Equal {
295                    // Apply ASC/DESC
296                    let asc = order_by.get(i).is_some_and(|o| matches!(o.direction, OrderDirection::Asc));
297                    return if asc { cmp } else { cmp.reverse() };
298                }
299            }
300            Ordering::Equal
301        });
302
303        // Move sorted rows back to results
304        results.extend(keyed_rows.into_iter().map(|(_, row)| row));
305
306        Ok(())
307    }
308
309    /// Apply LIMIT and OFFSET to results.
310    fn apply_arena_limit_offset(
311        &self,
312        mut results: Vec<Row>,
313        limit: Option<usize>,
314        offset: Option<usize>,
315    ) -> Vec<Row> {
316        // Apply offset first
317        if let Some(off) = offset {
318            if off >= results.len() {
319                return vec![];
320            }
321            results = results.into_iter().skip(off).collect();
322        }
323
324        // Apply limit
325        if let Some(lim) = limit {
326            results.truncate(lim);
327        }
328
329        results
330    }
331
332    /// Check if select list contains aggregate functions.
333    fn has_arena_aggregates<'arena>(&self, select_list: &[ArenaSelectItem<'arena>]) -> bool {
334        for item in select_list.iter() {
335            if let ArenaSelectItem::Expression { expr, .. } = item {
336                if self.arena_expr_has_aggregate(expr) {
337                    return true;
338                }
339            }
340        }
341        false
342    }
343
344    /// Check if an arena expression contains an aggregate function.
345    fn arena_expr_has_aggregate<'arena>(&self, expr: &ArenaExpression<'arena>) -> bool {
346        match expr {
347            // Hot-path inline variants
348            ArenaExpression::BinaryOp { left, right, .. } => {
349                self.arena_expr_has_aggregate(left) || self.arena_expr_has_aggregate(right)
350            }
351            ArenaExpression::UnaryOp { expr, .. } => self.arena_expr_has_aggregate(expr),
352            ArenaExpression::IsNull { expr, .. } => self.arena_expr_has_aggregate(expr),
353            ArenaExpression::Conjunction(children) | ArenaExpression::Disjunction(children) => {
354                children.iter().any(|c| self.arena_expr_has_aggregate(c))
355            }
356            ArenaExpression::Literal(_)
357            | ArenaExpression::Placeholder(_)
358            | ArenaExpression::NumberedPlaceholder(_)
359            | ArenaExpression::NamedPlaceholder(_)
360            | ArenaExpression::ColumnRef { .. }
361            | ArenaExpression::Wildcard
362            | ArenaExpression::CurrentDate
363            | ArenaExpression::CurrentTime { .. }
364            | ArenaExpression::CurrentTimestamp { .. }
365            | ArenaExpression::Default => false,
366            // Cold-path extended variants
367            ArenaExpression::Extended(ext) => self.arena_extended_has_aggregate(ext),
368        }
369    }
370
371    /// Check if an extended expression contains an aggregate function.
372    fn arena_extended_has_aggregate<'arena>(&self, ext: &ArenaExtendedExpr<'arena>) -> bool {
373        match ext {
374            ArenaExtendedExpr::AggregateFunction { .. } => true,
375            ArenaExtendedExpr::Function { args, .. } => {
376                args.iter().any(|a| self.arena_expr_has_aggregate(a))
377            }
378            ArenaExtendedExpr::Case { operand, when_clauses, else_result, .. } => {
379                operand.as_ref().is_some_and(|o| self.arena_expr_has_aggregate(o))
380                    || when_clauses.iter().any(|w| {
381                        w.conditions.iter().any(|c| self.arena_expr_has_aggregate(c))
382                            || self.arena_expr_has_aggregate(&w.result)
383                    })
384                    || else_result.as_ref().is_some_and(|e| self.arena_expr_has_aggregate(e))
385            }
386            ArenaExtendedExpr::Between { expr, low, high, .. } => {
387                self.arena_expr_has_aggregate(expr)
388                    || self.arena_expr_has_aggregate(low)
389                    || self.arena_expr_has_aggregate(high)
390            }
391            ArenaExtendedExpr::InList { expr, values, .. } => {
392                self.arena_expr_has_aggregate(expr)
393                    || values.iter().any(|v| self.arena_expr_has_aggregate(v))
394            }
395            ArenaExtendedExpr::Cast { expr, .. } => self.arena_expr_has_aggregate(expr),
396            ArenaExtendedExpr::Like { expr, pattern, .. } => {
397                self.arena_expr_has_aggregate(expr) || self.arena_expr_has_aggregate(pattern)
398            }
399            _ => false,
400        }
401    }
402}
vibesql_executor/select/executor/arena_execution.rs

vibesql_executor/select/executor/
arena_execution.rs
