radixdb_executor/aggregation/
storage.rs

1use super::*;
2
3impl<'host, H: AggregationHost + ?Sized> AggregationExecutor<'host, H> {
4    pub(super) fn storage_aggregate_expr_name(func_name: &str, fc: &FunctionCall) -> String {
5        if fc.arguments.is_empty() || matches!(fc.arguments.first(), Some(Expression::Star(_))) {
6            format!("{}(*)", func_name)
7        } else if let Some(Expression::Identifier(ident)) = fc.arguments.first() {
8            format!("{}({})", func_name, ident.value)
9        } else {
10            format!("{}(?)", func_name)
11        }
12    }
13
14    pub(super) fn storage_aggregate_from_call(
15        fc: &FunctionCall,
16        col_map: &FxHashMap<&str, usize>,
17    ) -> Option<(
18        String,
19        radixdb_storage::mvcc::version_store::AggregateOp,
20        usize,
21    )> {
22        use radixdb_storage::mvcc::version_store::AggregateOp;
23
24        if fc.filter.is_some() || fc.is_distinct || !fc.order_by.is_empty() {
25            return None;
26        }
27
28        let func_name = fc.function.to_uppercase();
29        let (op, col_idx) = match func_name.as_str() {
30            "COUNT" => {
31                if fc.arguments.is_empty()
32                    || matches!(fc.arguments.first(), Some(Expression::Star(_)))
33                {
34                    (AggregateOp::CountStar, 0)
35                } else if let Some(Expression::Identifier(ident)) = fc.arguments.first() {
36                    let col_name = ident.value.to_lowercase();
37                    (AggregateOp::Count, *col_map.get(col_name.as_str())?)
38                } else {
39                    return None;
40                }
41            }
42            "SUM" => {
43                if let Some(Expression::Identifier(ident)) = fc.arguments.first() {
44                    let col_name = ident.value.to_lowercase();
45                    (AggregateOp::Sum, *col_map.get(col_name.as_str())?)
46                } else {
47                    return None;
48                }
49            }
50            "AVG" => {
51                if let Some(Expression::Identifier(ident)) = fc.arguments.first() {
52                    let col_name = ident.value.to_lowercase();
53                    (AggregateOp::Avg, *col_map.get(col_name.as_str())?)
54                } else {
55                    return None;
56                }
57            }
58            "MIN" => {
59                if let Some(Expression::Identifier(ident)) = fc.arguments.first() {
60                    let col_name = ident.value.to_lowercase();
61                    (AggregateOp::Min, *col_map.get(col_name.as_str())?)
62                } else {
63                    return None;
64                }
65            }
66            "MAX" => {
67                if let Some(Expression::Identifier(ident)) = fc.arguments.first() {
68                    let col_name = ident.value.to_lowercase();
69                    (AggregateOp::Max, *col_map.get(col_name.as_str())?)
70                } else {
71                    return None;
72                }
73            }
74            _ => return None,
75        };
76
77        Some((
78            Self::storage_aggregate_expr_name(&func_name, fc),
79            op,
80            col_idx,
81        ))
82    }
83
84    pub(super) fn collect_storage_aggregate_dependencies(
85        expr: &Expression,
86        col_map: &FxHashMap<&str, usize>,
87        out: &mut Vec<(
88            String,
89            radixdb_storage::mvcc::version_store::AggregateOp,
90            usize,
91        )>,
92    ) -> bool {
93        match expr {
94            Expression::FunctionCall(fc) => {
95                if is_aggregate_function(&fc.function) {
96                    if let Some(dep) = Self::storage_aggregate_from_call(fc, col_map) {
97                        out.push(dep);
98                        true
99                    } else {
100                        false
101                    }
102                } else {
103                    fc.arguments
104                        .iter()
105                        .all(|arg| Self::collect_storage_aggregate_dependencies(arg, col_map, out))
106                        && fc.filter.as_ref().is_none_or(|filter| {
107                            Self::collect_storage_aggregate_dependencies(filter, col_map, out)
108                        })
109                        && fc.order_by.iter().all(|order| {
110                            Self::collect_storage_aggregate_dependencies(
111                                &order.expression,
112                                col_map,
113                                out,
114                            )
115                        })
116                }
117            }
118            Expression::Aliased(aliased) => {
119                Self::collect_storage_aggregate_dependencies(&aliased.expression, col_map, out)
120            }
121            Expression::Infix(infix) => {
122                Self::collect_storage_aggregate_dependencies(&infix.left, col_map, out)
123                    && Self::collect_storage_aggregate_dependencies(&infix.right, col_map, out)
124            }
125            Expression::Prefix(prefix) => {
126                Self::collect_storage_aggregate_dependencies(&prefix.right, col_map, out)
127            }
128            Expression::Distinct(distinct) => {
129                Self::collect_storage_aggregate_dependencies(&distinct.expr, col_map, out)
130            }
131            Expression::In(in_expr) => {
132                Self::collect_storage_aggregate_dependencies(&in_expr.left, col_map, out)
133                    && Self::collect_storage_aggregate_dependencies(&in_expr.right, col_map, out)
134            }
135            Expression::InHashSet(in_expr) => {
136                Self::collect_storage_aggregate_dependencies(&in_expr.column, col_map, out)
137            }
138            Expression::Between(between) => {
139                Self::collect_storage_aggregate_dependencies(&between.expr, col_map, out)
140                    && Self::collect_storage_aggregate_dependencies(&between.lower, col_map, out)
141                    && Self::collect_storage_aggregate_dependencies(&between.upper, col_map, out)
142            }
143            Expression::Like(like) => {
144                Self::collect_storage_aggregate_dependencies(&like.left, col_map, out)
145                    && Self::collect_storage_aggregate_dependencies(&like.pattern, col_map, out)
146                    && like.escape.as_ref().is_none_or(|escape| {
147                        Self::collect_storage_aggregate_dependencies(escape, col_map, out)
148                    })
149            }
150            Expression::List(list) => list
151                .elements
152                .iter()
153                .all(|item| Self::collect_storage_aggregate_dependencies(item, col_map, out)),
154            Expression::ExpressionList(list) => list
155                .expressions
156                .iter()
157                .all(|item| Self::collect_storage_aggregate_dependencies(item, col_map, out)),
158            Expression::Case(case) => {
159                case.value.as_ref().is_none_or(|value| {
160                    Self::collect_storage_aggregate_dependencies(value, col_map, out)
161                }) && case.when_clauses.iter().all(|when| {
162                    Self::collect_storage_aggregate_dependencies(&when.condition, col_map, out)
163                        && Self::collect_storage_aggregate_dependencies(
164                            &when.then_result,
165                            col_map,
166                            out,
167                        )
168                }) && case.else_value.as_ref().is_none_or(|else_value| {
169                    Self::collect_storage_aggregate_dependencies(else_value, col_map, out)
170                })
171            }
172            Expression::Cast(cast) => {
173                Self::collect_storage_aggregate_dependencies(&cast.expr, col_map, out)
174            }
175            Expression::AllAny(_)
176            | Expression::Exists(_)
177            | Expression::ScalarSubquery(_)
178            | Expression::Window(_)
179            | Expression::TableSource(_)
180            | Expression::JoinSource(_)
181            | Expression::SubquerySource(_)
182            | Expression::ValuesSource(_)
183            | Expression::CteReference(_)
184            | Expression::FunctionTableSource(_) => false,
185            Expression::Identifier(_)
186            | Expression::QualifiedIdentifier(_)
187            | Expression::IntegerLiteral(_)
188            | Expression::FloatLiteral(_)
189            | Expression::StringLiteral(_)
190            | Expression::BooleanLiteral(_)
191            | Expression::NullLiteral(_)
192            | Expression::IntervalLiteral(_)
193            | Expression::BoundValue(_)
194            | Expression::Parameter(_)
195            | Expression::Star(_)
196            | Expression::QualifiedStar(_)
197            | Expression::Default(_) => true,
198        }
199    }
200
201    /// Try to use storage-level aggregation for GROUP BY queries.
202    ///
203    /// This optimization bypasses row materialization by computing aggregates
204    /// directly from arena storage using Arc::clone for group keys.
205    ///
206    /// Returns None if the optimization cannot be applied.
207    ///
208    /// Currently only applies to simple queries with:
209    /// - GROUP BY columns that match SELECT identifiers exactly (same order)
210    /// - Simple aggregates (COUNT, SUM, AVG, MIN, MAX) on column references
211    /// - No ROLLUP, CUBE, or GROUPING SETS
212    /// - Optional WHERE if it can be fully converted to a storage expression
213    /// - Optional HAVING that can be evaluated against the grouped result row
214    pub fn try_storage_aggregation(
215        &self,
216        table: &dyn radixdb_storage::traits::Table,
217        stmt: &SelectStatement,
218        ctx: &ExecutionContext,
219        all_columns: &[String],
220        classification: &QueryClassification,
221    ) -> Option<Box<dyn QueryResult>> {
222        use radixdb_sql::ast::GroupByModifier;
223        use radixdb_storage::mvcc::version_store::AggregateOp;
224
225        // HAVING can be applied after the storage-level grouped rows are
226        // produced if it references columns present in that small result row.
227        if !classification.has_group_by {
228            return None;
229        }
230        if classification.where_has_parameters || classification.where_has_subqueries {
231            return None;
232        }
233
234        // Only for simple GROUP BY (no ROLLUP, CUBE, or GROUPING SETS)
235        if !matches!(stmt.group_by.modifier, GroupByModifier::None) {
236            return None;
237        }
238
239        // Only for simple GROUP BY expressions (column references)
240        let group_by_cols = &stmt.group_by.columns;
241        if group_by_cols.is_empty() {
242            return None;
243        }
244
245        // Build column name -> index map
246        let col_map: FxHashMap<&str, usize> = all_columns
247            .iter()
248            .enumerate()
249            .map(|(i, name)| (name.as_str(), i))
250            .collect();
251
252        // Extract group-by column names and indices (in GROUP BY order)
253        let mut group_by_indices: Vec<usize> = Vec::new();
254        let mut group_by_col_names: Vec<String> = Vec::new();
255        for expr in group_by_cols {
256            match expr {
257                Expression::Identifier(ident) => {
258                    let col_name = ident.value.to_lowercase().to_string();
259                    if let Some(&idx) = col_map.get(col_name.as_str()) {
260                        group_by_indices.push(idx);
261                        group_by_col_names.push(col_name);
262                    } else {
263                        return None; // Unknown column
264                    }
265                }
266                _ => return None, // Non-column GROUP BY not supported
267            }
268        }
269
270        // Parse SELECT columns - identify which are group-by columns vs aggregates
271        // Track positions so we can verify GROUP BY columns come first
272        let mut select_group_count = 0;
273        let mut seen_aggregate = false;
274        let mut aggregates: Vec<(AggregateOp, usize)> = Vec::new();
275        let mut agg_aliases: Vec<(String, usize)> = Vec::new();
276        let mut result_columns: Vec<String> = Vec::new();
277
278        for col_expr in &stmt.columns {
279            match col_expr {
280                Expression::Identifier(ident) => {
281                    // This must be a GROUP BY column
282                    let col_name_lower = ident.value.to_lowercase().to_string();
283                    if !group_by_col_names.contains(&col_name_lower) {
284                        return None; // Column not in GROUP BY
285                    }
286                    if seen_aggregate {
287                        // GROUP BY columns must come before aggregates for this optimization
288                        return None;
289                    }
290                    if group_by_col_names.get(select_group_count) != Some(&col_name_lower) {
291                        return None; // SELECT group columns must match GROUP BY order
292                    }
293                    select_group_count += 1;
294                    result_columns.push(ident.value.to_string());
295                }
296                Expression::FunctionCall(fc) => {
297                    seen_aggregate = true;
298                    let (expr_name, op, col_idx) = Self::storage_aggregate_from_call(fc, &col_map)?;
299                    let agg_idx = aggregates.len();
300                    agg_aliases.push((expr_name.clone(), group_by_indices.len() + agg_idx));
301                    aggregates.push((op, col_idx));
302
303                    result_columns.push(expr_name);
304                }
305                Expression::Aliased(aliased) => {
306                    // Handle aliased identifier (group by column with alias)
307                    if let Expression::Identifier(ident) = aliased.expression.as_ref() {
308                        let col_name_lower = ident.value.to_lowercase().to_string();
309                        if !group_by_col_names.contains(&col_name_lower) {
310                            return None; // Column not in GROUP BY
311                        }
312                        if seen_aggregate {
313                            return None;
314                        }
315                        if group_by_col_names.get(select_group_count) != Some(&col_name_lower) {
316                            return None; // SELECT group columns must match GROUP BY order
317                        }
318                        select_group_count += 1;
319                        result_columns.push(aliased.alias.value.to_string());
320                    }
321                    // Handle aliased aggregate
322                    else if let Expression::FunctionCall(fc) = aliased.expression.as_ref() {
323                        seen_aggregate = true;
324                        let (expr_name, op, col_idx) =
325                            Self::storage_aggregate_from_call(fc, &col_map)?;
326                        let agg_idx = aggregates.len();
327                        agg_aliases.push((expr_name, group_by_indices.len() + agg_idx));
328                        aggregates.push((op, col_idx));
329                        result_columns.push(aliased.alias.value.to_string());
330                    } else {
331                        return None; // Other aliased expression not supported
332                    }
333                }
334                _ => return None, // Unsupported expression type
335            }
336        }
337
338        // SELECT must include at least all GROUP BY columns (can have more)
339        // but for simplicity, require exact match with GROUP BY column count
340        if select_group_count != group_by_col_names.len() {
341            return None;
342        }
343
344        let public_column_count = result_columns.len();
345        let public_aggregate_count = aggregates.len();
346        for (idx, col_name) in group_by_col_names.iter().enumerate() {
347            if result_columns
348                .get(idx)
349                .is_some_and(|name| !name.eq_ignore_ascii_case(col_name))
350            {
351                agg_aliases.push((col_name.clone(), idx));
352            }
353        }
354        let mut known_agg_aliases: FxHashSet<String> = agg_aliases
355            .iter()
356            .map(|(name, _)| name.to_lowercase())
357            .collect();
358        let mut retained_order_dependencies = Vec::new();
359
360        // ORDER BY may reference an aggregate that is intentionally absent
361        // from the SELECT list. Keep direct hidden aggregates in the storage
362        // result until the outer ORDER BY has consumed them; the public
363        // projection removes them afterwards. More complex expressions still
364        // use the canonical aggregation path, which can rewrite aggregate
365        // subexpressions before evaluation.
366        for order_by in &stmt.order_by {
367            let mut hidden_deps = Vec::new();
368            if !Self::collect_storage_aggregate_dependencies(
369                &order_by.expression,
370                &col_map,
371                &mut hidden_deps,
372            ) {
373                return None;
374            }
375            let has_new_dependency = hidden_deps
376                .iter()
377                .any(|(name, _, _)| !known_agg_aliases.contains(&name.to_lowercase()));
378            if has_new_dependency
379                && !matches!(
380                    &order_by.expression,
381                    Expression::FunctionCall(function)
382                        if is_aggregate_function(&function.function)
383                )
384            {
385                return None;
386            }
387            for (expr_name, op, col_idx) in hidden_deps {
388                if !known_agg_aliases.insert(expr_name.to_lowercase()) {
389                    continue;
390                }
391                let agg_idx = aggregates.len();
392                agg_aliases.push((expr_name.clone(), group_by_indices.len() + agg_idx));
393                aggregates.push((op, col_idx));
394                retained_order_dependencies.push(expr_name);
395            }
396        }
397
398        if let Some(having) = stmt.having.as_ref() {
399            let mut hidden_deps = Vec::new();
400            if !Self::collect_storage_aggregate_dependencies(having, &col_map, &mut hidden_deps) {
401                return None;
402            }
403            for (expr_name, op, col_idx) in hidden_deps {
404                if !known_agg_aliases.insert(expr_name.to_lowercase()) {
405                    continue;
406                }
407                let agg_idx = aggregates.len();
408                agg_aliases.push((expr_name, group_by_indices.len() + agg_idx));
409                aggregates.push((op, col_idx));
410            }
411        }
412
413        let where_storage_expr = if let Some(where_expr) = stmt.where_clause.as_ref() {
414            let (storage_expr, needs_memory_filter) =
415                crate::pushdown::try_pushdown(where_expr, table.schema(), Some(ctx));
416            if needs_memory_filter {
417                return None;
418            }
419            Some(storage_expr?)
420        } else {
421            None
422        };
423
424        // Call storage-level aggregation
425        let results = if let Some(where_expr) = where_storage_expr.as_ref() {
426            table.compute_filtered_grouped_aggregates(
427                &group_by_indices,
428                &aggregates,
429                where_expr.as_ref(),
430            )?
431        } else {
432            table.compute_grouped_aggregates(&group_by_indices, &aggregates)?
433        };
434
435        // Convert to rows
436        let mut rows = RowVec::new();
437        for (row_id, r) in results.into_iter().enumerate() {
438            let mut values = r.group_values;
439            values.extend(r.aggregate_values);
440            rows.push((
441                row_id as i64,
442                Row::from_compact_vec(CompactVec::from_vec(values)),
443            ));
444        }
445
446        if let Some(having) = stmt.having.as_ref() {
447            let having_filter =
448                RowFilter::with_aliases_and_context(having, &result_columns, &agg_aliases, ctx)
449                    .ok()?;
450            let mut filtered_rows = RowVec::new();
451            for (_, row) in rows {
452                if having_filter.matches_checked(&row).ok()? {
453                    let row_id = filtered_rows.len() as i64;
454                    filtered_rows.push((row_id, row));
455                }
456            }
457            rows = filtered_rows;
458        }
459
460        let retained_column_count = public_column_count + retained_order_dependencies.len();
461        if aggregates.len() > public_aggregate_count + retained_order_dependencies.len() {
462            let mut projected_rows = RowVec::with_capacity(rows.len());
463            for (_, row) in rows {
464                let mut values = CompactVec::with_capacity(retained_column_count);
465                for value in row.as_slice().iter().take(retained_column_count) {
466                    values.push(value.clone());
467                }
468                let row_id = projected_rows.len() as i64;
469                projected_rows.push((row_id, Row::from_compact_vec(values)));
470            }
471            rows = projected_rows;
472        }
473
474        result_columns.extend(retained_order_dependencies);
475        result_columns.truncate(retained_column_count);
476        Some(Box::new(ExecutorResult::new(result_columns, rows)))
477    }
478
479    /// Try fast COUNT(DISTINCT col) using compiled cache
480    ///
481    /// This is a compiled fast path for simple `SELECT COUNT(DISTINCT col) FROM table` queries.
482    /// On first execution, it analyzes and compiles the query pattern.
483    /// On subsequent executions, it skips parsing and directly fetches the distinct count.
484    ///
485    /// # Arguments
486    /// * `stmt` - The SELECT statement
487    /// * `compiled` - The compiled execution state (shared via RwLock)
488    ///
489    /// # Returns
490    /// * `Some(Ok(result))` - Query succeeded via fast path
491    /// * `Some(Err(e))` - Query failed
492    /// * `None` - Query doesn't qualify for this fast path (use normal path)
493    pub(crate) fn try_fast_count_distinct_compiled(
494        &self,
495        stmt: &SelectStatement,
496        compiled: &RwLock<CompiledExecution>,
497    ) -> Option<Result<Box<dyn QueryResult>>> {
498        // Quick reject: must not be in an explicit transaction (for simplicity)
499        {
500            let active_tx = match self.host.aggregation_active_transaction().try_lock() {
501                Ok(guard) => guard,
502                Err(_) => return None,
503            };
504            if active_tx.is_some() {
505                return None;
506            }
507        }
508
509        // Try read lock first - check if already compiled
510        {
511            let compiled_guard = match compiled.read() {
512                Ok(guard) => guard,
513                Err(_) => return None,
514            };
515            match &*compiled_guard {
516                CompiledExecution::NotOptimizable(epoch)
517                    if self.host.aggregation_engine().schema_epoch() == *epoch =>
518                {
519                    return None
520                }
521                CompiledExecution::CountDistinct(cd) => {
522                    // Fast path: validate epoch and execute directly
523                    if self.host.aggregation_engine().schema_epoch() == cd.cached_epoch {
524                        return Some(self.execute_compiled_count_distinct(cd));
525                    }
526                    // Schema changed - fall through to recompile
527                }
528                CompiledExecution::NotOptimizable(_) | CompiledExecution::Unknown => {} // Epoch changed or first run - fall through to recompile
529                // Other variants - not a COUNT DISTINCT query
530                _ => return None,
531            }
532        }
533
534        // First execution or schema changed - compile and cache (write lock)
535        self.compile_and_execute_count_distinct(stmt, compiled)
536    }
537
538    /// Execute using pre-compiled COUNT(DISTINCT col) info
539    pub(super) fn execute_compiled_count_distinct(
540        &self,
541        cd: &CompiledCountDistinct,
542    ) -> Result<Box<dyn QueryResult>> {
543        // Get table and count distinct values directly
544        let tx = self.host.aggregation_engine().begin_transaction()?;
545        let table = tx.get_table(&cd.table_name)?;
546
547        let count = table
548            .get_partition_count(&cd.column_name)
549            .ok_or_else(|| radixdb_core::Error::internal("Index no longer available for column"))?;
550
551        // Build result
552        let mut result_values = CompactVec::with_capacity(1);
553        result_values.push(Value::Integer(count as i64));
554        let row = Row::from_compact_vec(result_values);
555        let mut rows = RowVec::with_capacity(1);
556        rows.push((0, row));
557
558        Ok(Box::new(ExecutorResult::new(
559            vec![cd.result_column_name.clone()],
560            rows,
561        )))
562    }
563
564    /// Compile and execute COUNT(DISTINCT col), caching the compiled state
565    pub(super) fn compile_and_execute_count_distinct(
566        &self,
567        stmt: &SelectStatement,
568        compiled: &RwLock<CompiledExecution>,
569    ) -> Option<Result<Box<dyn QueryResult>>> {
570        use radixdb_core::SmartString;
571
572        // Acquire write lock
573        let mut compiled_guard = match compiled.write() {
574            Ok(guard) => guard,
575            Err(_) => return None,
576        };
577
578        // Double-check (another thread may have compiled while we waited)
579        match &*compiled_guard {
580            CompiledExecution::NotOptimizable(epoch)
581                if self.host.aggregation_engine().schema_epoch() == *epoch =>
582            {
583                return None
584            }
585            CompiledExecution::CountDistinct(cd) => {
586                if self.host.aggregation_engine().schema_epoch() == cd.cached_epoch {
587                    return Some(self.execute_compiled_count_distinct(cd));
588                }
589                // Schema changed, continue to recompile
590            }
591            CompiledExecution::NotOptimizable(_) | CompiledExecution::Unknown => {} // Epoch changed or first run - recompile
592            _ => return None,
593        }
594
595        // Pattern detection: SELECT COUNT(DISTINCT col) FROM table
596        // Must have:
597        // - Exactly one column expression
598        // - That column is COUNT(DISTINCT col) function call
599        // - No WHERE, GROUP BY, HAVING, ORDER BY, LIMIT, CTEs, set operations
600        // - Single table source (no joins)
601
602        if stmt.columns.len() != 1 {
603            *compiled_guard =
604                CompiledExecution::NotOptimizable(self.host.aggregation_engine().schema_epoch());
605            return None;
606        }
607
608        // Check for disqualifying clauses
609        if stmt.where_clause.is_some()
610            || !stmt.group_by.columns.is_empty()
611            || stmt.having.is_some()
612            || !stmt.order_by.is_empty()
613            || stmt.limit.is_some()
614            || stmt.offset.is_some()
615            || stmt.with.is_some()
616            || !stmt.set_operations.is_empty()
617        {
618            *compiled_guard =
619                CompiledExecution::NotOptimizable(self.host.aggregation_engine().schema_epoch());
620            return None;
621        }
622
623        // Extract COUNT(DISTINCT col) pattern
624        let (column_name, result_column_name) = match &stmt.columns[0] {
625            Expression::FunctionCall(func) => {
626                // Must be COUNT function
627                if func.function.to_uppercase() != "COUNT" {
628                    *compiled_guard = CompiledExecution::NotOptimizable(
629                        self.host.aggregation_engine().schema_epoch(),
630                    );
631                    return None;
632                }
633                // Must be DISTINCT - if not, don't mark as NotOptimizable
634                // because COUNT(*) fast path might handle it
635                if !func.is_distinct {
636                    return None;
637                }
638                if func.arguments.len() != 1 {
639                    *compiled_guard = CompiledExecution::NotOptimizable(
640                        self.host.aggregation_engine().schema_epoch(),
641                    );
642                    return None;
643                }
644                // Get column name from argument
645                let col = match &func.arguments[0] {
646                    Expression::Identifier(ident) => ident.value.to_lowercase(),
647                    _ => {
648                        *compiled_guard = CompiledExecution::NotOptimizable(
649                            self.host.aggregation_engine().schema_epoch(),
650                        );
651                        return None;
652                    }
653                };
654                let result_name = format!("COUNT(DISTINCT {})", col);
655                (col, result_name)
656            }
657            Expression::Aliased(aliased) => {
658                // Handle COUNT(DISTINCT col) AS alias
659                match aliased.expression.as_ref() {
660                    Expression::FunctionCall(func) => {
661                        // Must be COUNT function
662                        if func.function.to_uppercase() != "COUNT" {
663                            *compiled_guard = CompiledExecution::NotOptimizable(
664                                self.host.aggregation_engine().schema_epoch(),
665                            );
666                            return None;
667                        }
668                        // Must be DISTINCT - if not, don't mark as NotOptimizable
669                        // because COUNT(*) fast path might handle it
670                        if !func.is_distinct {
671                            return None;
672                        }
673                        if func.arguments.len() != 1 {
674                            *compiled_guard = CompiledExecution::NotOptimizable(
675                                self.host.aggregation_engine().schema_epoch(),
676                            );
677                            return None;
678                        }
679                        let col = match &func.arguments[0] {
680                            Expression::Identifier(ident) => ident.value.to_lowercase(),
681                            _ => {
682                                *compiled_guard = CompiledExecution::NotOptimizable(
683                                    self.host.aggregation_engine().schema_epoch(),
684                                );
685                                return None;
686                            }
687                        };
688                        (col, aliased.alias.value.to_string())
689                    }
690                    _ => {
691                        *compiled_guard = CompiledExecution::NotOptimizable(
692                            self.host.aggregation_engine().schema_epoch(),
693                        );
694                        return None;
695                    }
696                }
697            }
698            _ => {
699                *compiled_guard = CompiledExecution::NotOptimizable(
700                    self.host.aggregation_engine().schema_epoch(),
701                );
702                return None;
703            }
704        };
705
706        // Extract table name (bail out for AS OF temporal queries)
707        let table_name = match stmt.table_expr.as_deref() {
708            Some(Expression::TableSource(ts)) => {
709                if ts.as_of.is_some() {
710                    *compiled_guard = CompiledExecution::NotOptimizable(
711                        self.host.aggregation_engine().schema_epoch(),
712                    );
713                    return None;
714                }
715                ts.name.value_lower.clone()
716            }
717            _ => {
718                *compiled_guard = CompiledExecution::NotOptimizable(
719                    self.host.aggregation_engine().schema_epoch(),
720                );
721                return None;
722            }
723        };
724
725        // Try to get table and verify index exists
726        let tx = match self.host.aggregation_engine().begin_transaction() {
727            Ok(tx) => tx,
728            Err(_) => {
729                *compiled_guard = CompiledExecution::NotOptimizable(
730                    self.host.aggregation_engine().schema_epoch(),
731                );
732                return None;
733            }
734        };
735
736        let table = match tx.get_table(&table_name) {
737            Ok(t) => t,
738            Err(_) => {
739                *compiled_guard = CompiledExecution::NotOptimizable(
740                    self.host.aggregation_engine().schema_epoch(),
741                );
742                return None;
743            }
744        };
745
746        // Check if column has an index (required for fast path)
747        if table.get_partition_count(&column_name).is_none() {
748            *compiled_guard =
749                CompiledExecution::NotOptimizable(self.host.aggregation_engine().schema_epoch());
750            return None;
751        }
752
753        // Get the count
754        let count = table.get_partition_count(&column_name).unwrap();
755
756        // Cache the compiled state
757        let compiled_cd = CompiledCountDistinct {
758            table_name: SmartString::new(&table_name),
759            column_name: SmartString::new(&column_name),
760            result_column_name: result_column_name.clone(),
761            cached_epoch: self.host.aggregation_engine().schema_epoch(),
762        };
763        *compiled_guard = CompiledExecution::CountDistinct(compiled_cd);
764        drop(compiled_guard);
765
766        // Build result
767        let mut result_values = CompactVec::with_capacity(1);
768        result_values.push(Value::Integer(count as i64));
769        let row = Row::from_compact_vec(result_values);
770        let mut rows = RowVec::with_capacity(1);
771        rows.push((0, row));
772
773        Some(Ok(Box::new(ExecutorResult::new(
774            vec![result_column_name],
775            rows,
776        ))))
777    }
778
779    /// COUNT(*) fast path for simple queries
780    ///
781    /// This is a compiled fast path for simple `SELECT COUNT(*) FROM table` queries.
782    /// On first execution, it analyzes and compiles the query pattern.
783    /// On subsequent executions, it skips parsing and directly fetches the row count.
784    ///
785    /// # Arguments
786    /// * `stmt` - The SELECT statement
787    /// * `compiled` - The compiled execution state (shared via RwLock)
788    ///
789    /// # Returns
790    /// * `Some(Ok(result))` - Query succeeded via fast path
791    /// * `Some(Err(e))` - Query failed
792    /// * `None` - Query doesn't qualify for this fast path (use normal path)
793    pub(crate) fn try_fast_count_star_compiled(
794        &self,
795        stmt: &SelectStatement,
796        compiled: &RwLock<CompiledExecution>,
797    ) -> Option<Result<Box<dyn QueryResult>>> {
798        // Quick reject: must not be in an explicit transaction (for simplicity)
799        {
800            let active_tx = match self.host.aggregation_active_transaction().try_lock() {
801                Ok(guard) => guard,
802                Err(_) => return None,
803            };
804            if active_tx.is_some() {
805                return None;
806            }
807        }
808
809        // Try read lock first - check if already compiled
810        {
811            let compiled_guard = match compiled.read() {
812                Ok(guard) => guard,
813                Err(_) => return None,
814            };
815            match &*compiled_guard {
816                CompiledExecution::NotOptimizable(epoch)
817                    if self.host.aggregation_engine().schema_epoch() == *epoch =>
818                {
819                    return None
820                }
821                CompiledExecution::CountStar(cs) => {
822                    // Fast path: validate epoch and execute directly
823                    if self.host.aggregation_engine().schema_epoch() == cs.cached_epoch {
824                        return Some(self.execute_compiled_count_star(cs));
825                    }
826                    // Schema changed - fall through to recompile
827                }
828                CompiledExecution::NotOptimizable(_) | CompiledExecution::Unknown => {} // Epoch changed or first run - fall through to recompile
829                // Other variants - not a COUNT(*) query
830                _ => return None,
831            }
832        }
833
834        // First execution or schema changed - compile and cache (write lock)
835        self.compile_and_execute_count_star(stmt, compiled)
836    }
837
838    /// Execute using pre-compiled COUNT(*) info
839    pub(super) fn execute_compiled_count_star(
840        &self,
841        cs: &crate::compiled_plan::CompiledCountStar,
842    ) -> Result<Box<dyn QueryResult>> {
843        // Get table and count rows directly
844        let tx = self.host.aggregation_engine().begin_transaction()?;
845        let table = tx.get_table(&cs.table_name)?;
846
847        let count = table.row_count();
848
849        // Build result
850        let mut result_values = CompactVec::with_capacity(1);
851        result_values.push(Value::Integer(count as i64));
852        let row = Row::from_compact_vec(result_values);
853        let mut rows = RowVec::with_capacity(1);
854        rows.push((0, row));
855
856        Ok(Box::new(ExecutorResult::new(
857            vec![cs.result_column_name.clone()],
858            rows,
859        )))
860    }
861
862    /// Compile and execute COUNT(*), caching the compiled state
863    pub(super) fn compile_and_execute_count_star(
864        &self,
865        stmt: &SelectStatement,
866        compiled: &RwLock<CompiledExecution>,
867    ) -> Option<Result<Box<dyn QueryResult>>> {
868        use crate::compiled_plan::CompiledCountStar;
869        use radixdb_core::SmartString;
870
871        // Acquire write lock
872        let mut compiled_guard = match compiled.write() {
873            Ok(guard) => guard,
874            Err(_) => return None,
875        };
876
877        // Double-check (another thread may have compiled while we waited)
878        match &*compiled_guard {
879            CompiledExecution::NotOptimizable(epoch)
880                if self.host.aggregation_engine().schema_epoch() == *epoch =>
881            {
882                return None
883            }
884            CompiledExecution::CountStar(cs) => {
885                if self.host.aggregation_engine().schema_epoch() == cs.cached_epoch {
886                    return Some(self.execute_compiled_count_star(cs));
887                }
888                // Schema changed, continue to recompile
889            }
890            CompiledExecution::NotOptimizable(_) | CompiledExecution::Unknown => {} // Epoch changed or first run - recompile
891            _ => return None,
892        }
893
894        // Pattern detection: SELECT COUNT(*) FROM table
895        // Must have:
896        // - Exactly one column expression
897        // - That column is COUNT(*) or COUNT(1) function call (not DISTINCT)
898        // - No WHERE, GROUP BY, HAVING, ORDER BY, LIMIT, CTEs, set operations
899        // - Single table source (no joins)
900
901        if stmt.columns.len() != 1 {
902            *compiled_guard =
903                CompiledExecution::NotOptimizable(self.host.aggregation_engine().schema_epoch());
904            return None;
905        }
906
907        // Check for disqualifying clauses
908        if stmt.where_clause.is_some()
909            || !stmt.group_by.columns.is_empty()
910            || stmt.having.is_some()
911            || !stmt.order_by.is_empty()
912            || stmt.limit.is_some()
913            || stmt.offset.is_some()
914            || stmt.with.is_some()
915            || !stmt.set_operations.is_empty()
916        {
917            *compiled_guard =
918                CompiledExecution::NotOptimizable(self.host.aggregation_engine().schema_epoch());
919            return None;
920        }
921
922        // Extract COUNT(*) pattern
923        let result_column_name = match &stmt.columns[0] {
924            Expression::FunctionCall(func) => {
925                // Must be COUNT function
926                if func.function.to_uppercase() != "COUNT" {
927                    *compiled_guard = CompiledExecution::NotOptimizable(
928                        self.host.aggregation_engine().schema_epoch(),
929                    );
930                    return None;
931                }
932                // Must NOT be DISTINCT - if DISTINCT, don't mark as NotOptimizable
933                // because COUNT DISTINCT fast path might handle it
934                if func.is_distinct {
935                    return None;
936                }
937                // Must not have FILTER clause
938                if func.filter.is_some() {
939                    *compiled_guard = CompiledExecution::NotOptimizable(
940                        self.host.aggregation_engine().schema_epoch(),
941                    );
942                    return None;
943                }
944                // Must be COUNT(*) or COUNT(1)
945                match func.arguments.len() {
946                    0 => {
947                        // COUNT(*) without explicit star is rare but handle it
948                        "COUNT(*)".to_string()
949                    }
950                    1 => {
951                        match &func.arguments[0] {
952                            Expression::Star(_) => "COUNT(*)".to_string(),
953                            Expression::IntegerLiteral(lit) => {
954                                // COUNT(1) is equivalent to COUNT(*)
955                                if lit.value == 1 {
956                                    "COUNT(1)".to_string()
957                                } else {
958                                    *compiled_guard = CompiledExecution::NotOptimizable(
959                                        self.host.aggregation_engine().schema_epoch(),
960                                    );
961                                    return None;
962                                }
963                            }
964                            _ => {
965                                // COUNT(col) without DISTINCT - not our fast path
966                                *compiled_guard = CompiledExecution::NotOptimizable(
967                                    self.host.aggregation_engine().schema_epoch(),
968                                );
969                                return None;
970                            }
971                        }
972                    }
973                    _ => {
974                        *compiled_guard = CompiledExecution::NotOptimizable(
975                            self.host.aggregation_engine().schema_epoch(),
976                        );
977                        return None;
978                    }
979                }
980            }
981            Expression::Aliased(aliased) => {
982                // Handle COUNT(*) AS alias
983                match aliased.expression.as_ref() {
984                    Expression::FunctionCall(func) => {
985                        // Must be COUNT function
986                        if func.function.to_uppercase() != "COUNT" {
987                            *compiled_guard = CompiledExecution::NotOptimizable(
988                                self.host.aggregation_engine().schema_epoch(),
989                            );
990                            return None;
991                        }
992                        // Must NOT be DISTINCT - if DISTINCT, don't mark as NotOptimizable
993                        // because COUNT DISTINCT fast path might handle it
994                        if func.is_distinct {
995                            return None;
996                        }
997                        // Must not have FILTER clause
998                        if func.filter.is_some() {
999                            *compiled_guard = CompiledExecution::NotOptimizable(
1000                                self.host.aggregation_engine().schema_epoch(),
1001                            );
1002                            return None;
1003                        }
1004                        match func.arguments.len() {
1005                            0 => aliased.alias.value.to_string(),
1006                            1 => match &func.arguments[0] {
1007                                Expression::Star(_) => aliased.alias.value.to_string(),
1008                                Expression::IntegerLiteral(lit) if lit.value == 1 => {
1009                                    aliased.alias.value.to_string()
1010                                }
1011                                Expression::IntegerLiteral(_) => {
1012                                    *compiled_guard = CompiledExecution::NotOptimizable(
1013                                        self.host.aggregation_engine().schema_epoch(),
1014                                    );
1015                                    return None;
1016                                }
1017                                _ => {
1018                                    *compiled_guard = CompiledExecution::NotOptimizable(
1019                                        self.host.aggregation_engine().schema_epoch(),
1020                                    );
1021                                    return None;
1022                                }
1023                            },
1024                            _ => {
1025                                *compiled_guard = CompiledExecution::NotOptimizable(
1026                                    self.host.aggregation_engine().schema_epoch(),
1027                                );
1028                                return None;
1029                            }
1030                        }
1031                    }
1032                    _ => {
1033                        *compiled_guard = CompiledExecution::NotOptimizable(
1034                            self.host.aggregation_engine().schema_epoch(),
1035                        );
1036                        return None;
1037                    }
1038                }
1039            }
1040            _ => {
1041                *compiled_guard = CompiledExecution::NotOptimizable(
1042                    self.host.aggregation_engine().schema_epoch(),
1043                );
1044                return None;
1045            }
1046        };
1047
1048        // Extract table name (bail out for AS OF temporal queries)
1049        let table_name = match stmt.table_expr.as_deref() {
1050            Some(Expression::TableSource(ts)) => {
1051                if ts.as_of.is_some() {
1052                    *compiled_guard = CompiledExecution::NotOptimizable(
1053                        self.host.aggregation_engine().schema_epoch(),
1054                    );
1055                    return None;
1056                }
1057                ts.name.value_lower.clone()
1058            }
1059            _ => {
1060                *compiled_guard = CompiledExecution::NotOptimizable(
1061                    self.host.aggregation_engine().schema_epoch(),
1062                );
1063                return None;
1064            }
1065        };
1066
1067        // Try to get table and get count
1068        let tx = match self.host.aggregation_engine().begin_transaction() {
1069            Ok(tx) => tx,
1070            Err(_) => {
1071                *compiled_guard = CompiledExecution::NotOptimizable(
1072                    self.host.aggregation_engine().schema_epoch(),
1073                );
1074                return None;
1075            }
1076        };
1077
1078        let table = match tx.get_table(&table_name) {
1079            Ok(t) => t,
1080            Err(_) => {
1081                *compiled_guard = CompiledExecution::NotOptimizable(
1082                    self.host.aggregation_engine().schema_epoch(),
1083                );
1084                return None;
1085            }
1086        };
1087
1088        // Get the count
1089        let count = table.row_count();
1090
1091        // Cache the compiled state
1092        let compiled_cs = CompiledCountStar {
1093            table_name: SmartString::new(&table_name),
1094            result_column_name: result_column_name.clone(),
1095            cached_epoch: self.host.aggregation_engine().schema_epoch(),
1096        };
1097        *compiled_guard = CompiledExecution::CountStar(compiled_cs);
1098        drop(compiled_guard);
1099
1100        // Build result
1101        let mut result_values = CompactVec::with_capacity(1);
1102        result_values.push(Value::Integer(count as i64));
1103        let row = Row::from_compact_vec(result_values);
1104        let mut rows = RowVec::with_capacity(1);
1105        rows.push((0, row));
1106
1107        Some(Ok(Box::new(ExecutorResult::new(
1108            vec![result_column_name],
1109            rows,
1110        ))))
1111    }
1112}
radixdb_executor/aggregation/storage.rs

radixdb_executor/aggregation/
storage.rs
