vibesql_executor/select/columnar/aggregate/
mod.rs

1//! Columnar aggregation - high-performance aggregate computation
2//!
3//! This module provides efficient aggregate operations on columnar data,
4//! organized into several sub-modules:
5//!
6//! - [`functions`] - Core aggregate function implementations (SUM, COUNT, AVG, MIN, MAX)
7//! - [`expression`] - Expression aggregates (e.g., SUM(a * b))
8//! - [`group_by`] - Hash-based GROUP BY aggregation
9//!
10//! The public API maintains backward compatibility while the implementation
11//! is split across focused modules for better maintainability.
12
13mod expression;
14mod functions;
15mod group_by;
16
17// Re-export public types and functions to maintain API compatibility
18pub use expression::{
19    evaluate_expression_to_column, evaluate_expression_with_cached_column, extract_aggregates,
20};
21pub use group_by::columnar_group_by;
22// SIMD-accelerated GROUP BY for ColumnarBatch - used in native columnar execution path
23pub use group_by::columnar_group_by_batch;
24use vibesql_ast::Expression;
25use vibesql_storage::Row;
26use vibesql_types::SqlValue;
27
28use super::{batch::ColumnarBatch, scan::ColumnarScan};
29use crate::{errors::ExecutorError, schema::CombinedSchema};
30
31/// Aggregate operation type
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum AggregateOp {
34    Sum,
35    Count,
36    Avg,
37    Min,
38    Max,
39}
40
41/// Source of data for an aggregate - either a simple column or an expression
42#[derive(Debug, Clone, PartialEq)]
43pub enum AggregateSource {
44    /// Simple column reference (fast path) - just a column index
45    Column(usize),
46    /// Complex expression that needs evaluation (e.g., a * b)
47    Expression(Expression),
48    /// COUNT(*) - count all rows (not tied to any specific column)
49    CountStar,
50}
51
52/// A complete aggregate specification
53#[derive(Debug, Clone)]
54pub struct AggregateSpec {
55    pub op: AggregateOp,
56    pub source: AggregateSource,
57}
58
59/// Compute an aggregate over a column with optional filtering
60///
61/// This is the core columnar aggregation function that processes
62/// columns directly without materializing Row objects.
63///
64/// Automatically detects when SIMD optimization is available (Int64/Float64 columns)
65/// and falls back to scalar implementation for other types.
66///
67/// # Arguments
68///
69/// * `scan` - Columnar scan over the data
70/// * `column_idx` - Index of the column to aggregate
71/// * `op` - Aggregate operation (SUM, COUNT, etc.)
72/// * `filter_bitmap` - Optional bitmap of which rows to include
73///
74/// # Returns
75///
76/// The aggregated SqlValue
77pub fn compute_columnar_aggregate(
78    scan: &ColumnarScan,
79    column_idx: usize,
80    op: AggregateOp,
81    filter_bitmap: Option<&[bool]>,
82) -> Result<SqlValue, ExecutorError> {
83    // Try SIMD path for numeric columns (5-10x speedup via LLVM auto-vectorization)
84    {
85        use super::simd_aggregate::{
86            can_use_simd_for_column, simd_aggregate_f64, simd_aggregate_i64,
87        };
88
89        // Detect if column is SIMD-compatible
90        if let Some(is_integer) = can_use_simd_for_column(scan, column_idx) {
91            // Use SIMD implementation for Int64/Float64 columns
92            return if is_integer {
93                simd_aggregate_i64(scan, column_idx, op, filter_bitmap)
94            } else {
95                simd_aggregate_f64(scan, column_idx, op, filter_bitmap)
96            };
97        }
98        // Fall through to scalar path for non-SIMD types
99    }
100
101    // Scalar fallback path (always available, used for String, Date, etc.)
102    functions::compute_columnar_aggregate_impl(scan, column_idx, op, filter_bitmap)
103}
104
105/// Compute multiple aggregates in a single pass over the data
106///
107/// This is more efficient than computing each aggregate separately
108/// as it only scans the data once.
109pub fn compute_multiple_aggregates(
110    rows: &[Row],
111    aggregates: &[AggregateSpec],
112    filter_bitmap: Option<&[bool]>,
113    schema: Option<&CombinedSchema>,
114) -> Result<Vec<SqlValue>, ExecutorError> {
115    let scan = ColumnarScan::new(rows);
116    let mut results = Vec::with_capacity(aggregates.len());
117
118    for spec in aggregates {
119        let result = match &spec.source {
120            // Fast path: direct column aggregation
121            AggregateSource::Column(column_idx) => {
122                compute_columnar_aggregate(&scan, *column_idx, spec.op, filter_bitmap)?
123            }
124            // Expression path: evaluate expression for each row, then aggregate
125            AggregateSource::Expression(expr) => {
126                let schema = schema.ok_or_else(|| {
127                    ExecutorError::UnsupportedExpression(
128                        "Schema required for expression aggregates".to_string(),
129                    )
130                })?;
131                expression::compute_expression_aggregate(
132                    rows,
133                    expr,
134                    spec.op,
135                    filter_bitmap,
136                    schema,
137                )?
138            }
139            // COUNT(*) path: count all rows
140            AggregateSource::CountStar => functions::compute_count(&scan, filter_bitmap)?,
141        };
142        results.push(result);
143    }
144
145    Ok(results)
146}
147
148/// Compute aggregates directly from a ColumnarBatch (no row conversion)
149///
150/// This is a high-performance path that eliminates the overhead of converting
151/// ColumnarBatch back to rows before aggregation. It works directly on the
152/// typed column arrays for maximum efficiency.
153///
154/// # Arguments
155///
156/// * `batch` - The ColumnarBatch to aggregate over (typically filtered)
157/// * `aggregates` - List of aggregate specifications to compute
158/// * `schema` - Optional schema for expression aggregates
159///
160/// # Performance
161///
162/// This function provides 20-30% speedup over the row-based path by:
163/// - Avoiding `batch.to_rows()` conversion overhead (~10-15ms for large batches)
164/// - Working directly on typed arrays (Int64, Float64) with SIMD operations
165/// - Reducing memory allocations
166///
167/// # Returns
168///
169/// Vector of aggregate results in the same order as the input aggregates
170pub fn compute_aggregates_from_batch(
171    batch: &ColumnarBatch,
172    aggregates: &[AggregateSpec],
173    schema: Option<&CombinedSchema>,
174) -> Result<Vec<SqlValue>, ExecutorError> {
175    // Handle empty batch
176    if batch.row_count() == 0 {
177        return Ok(aggregates
178            .iter()
179            .map(|spec| match spec.op {
180                AggregateOp::Count => SqlValue::Integer(0),
181                _ => SqlValue::Null,
182            })
183            .collect());
184    }
185
186    let mut results = Vec::with_capacity(aggregates.len());
187
188    for spec in aggregates {
189        let result = match &spec.source {
190            // Fast path: direct batch aggregation (no row conversion)
191            AggregateSource::Column(column_idx) => {
192                functions::compute_batch_aggregate(batch, *column_idx, spec.op)?
193            }
194            // Expression path: SIMD-accelerated evaluation directly on batch columns
195            // This eliminates the ~10-15ms overhead of batch.to_rows() for large batches
196            AggregateSource::Expression(expr) => {
197                let schema = schema.ok_or_else(|| {
198                    ExecutorError::UnsupportedExpression(
199                        "Schema required for expression aggregates".to_string(),
200                    )
201                })?;
202                // Try batch-native expression evaluation first
203                // Falls back to row-based evaluation if column types aren't supported
204                match expression::compute_batch_expression_aggregate(batch, expr, spec.op, schema) {
205                    Ok(value) => value,
206                    Err(ExecutorError::UnsupportedExpression(_)) => {
207                        // Fall back to row-based for unsupported column types (Mixed, Date, etc.)
208                        let rows = batch.to_rows()?;
209                        expression::compute_expression_aggregate(
210                            &rows, expr, spec.op, None, schema,
211                        )?
212                    }
213                    Err(other) => return Err(other),
214                }
215            }
216            // COUNT(*) path: just count rows in batch
217            AggregateSource::CountStar => SqlValue::Integer(batch.row_count() as i64),
218        };
219        results.push(result);
220    }
221
222    Ok(results)
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    fn make_test_rows() -> Vec<Row> {
230        vec![
231            Row::new(vec![SqlValue::Integer(10), SqlValue::Double(1.5)]),
232            Row::new(vec![SqlValue::Integer(20), SqlValue::Double(2.5)]),
233            Row::new(vec![SqlValue::Integer(30), SqlValue::Double(3.5)]),
234        ]
235    }
236
237    #[test]
238    fn test_sum_aggregate() {
239        let rows = make_test_rows();
240        let scan = ColumnarScan::new(&rows);
241
242        let result = functions::compute_sum(&scan, 0, None).unwrap();
243        assert_eq!(result, SqlValue::Integer(60));
244
245        let result = functions::compute_sum(&scan, 1, None).unwrap();
246        assert!(matches!(result, SqlValue::Double(sum) if (sum - 7.5).abs() < 0.001));
247    }
248
249    #[test]
250    fn test_count_aggregate() {
251        let rows = make_test_rows();
252        let scan = ColumnarScan::new(&rows);
253
254        let result = functions::compute_count(&scan, None).unwrap();
255        assert_eq!(result, SqlValue::Integer(3));
256    }
257
258    #[test]
259    fn test_sum_with_filter() {
260        let rows = make_test_rows();
261        let scan = ColumnarScan::new(&rows);
262        let filter = vec![true, false, true]; // Include rows 0 and 2
263
264        let result = functions::compute_sum(&scan, 0, Some(&filter)).unwrap();
265        assert_eq!(result, SqlValue::Integer(40));
266    }
267
268    #[test]
269    fn test_multiple_aggregates() {
270        let rows = make_test_rows();
271        let aggregates = vec![
272            AggregateSpec { op: AggregateOp::Sum, source: AggregateSource::Column(0) },
273            AggregateSpec { op: AggregateOp::Avg, source: AggregateSource::Column(1) },
274        ];
275
276        let results = compute_multiple_aggregates(&rows, &aggregates, None, None).unwrap();
277        assert_eq!(results.len(), 2);
278        assert_eq!(results[0], SqlValue::Integer(60));
279        assert!(matches!(results[1], SqlValue::Double(avg) if (avg - 2.5).abs() < 0.001));
280    }
281
282    #[test]
283    fn test_extract_aggregates_simple() {
284        use vibesql_catalog::{ColumnSchema, TableSchema};
285        use vibesql_types::DataType;
286
287        use crate::schema::CombinedSchema;
288
289        // Create a simple schema with two columns
290        let schema = TableSchema::new(
291            "test".to_string(),
292            vec![
293                ColumnSchema::new("col1".to_string(), DataType::Integer, false),
294                ColumnSchema::new("col2".to_string(), DataType::DoublePrecision, false),
295            ],
296        );
297
298        let combined_schema = CombinedSchema::from_table("test".to_string(), schema);
299
300        // Test SUM(col1)
301        let exprs = vec![Expression::AggregateFunction {
302            name: vibesql_ast::FunctionIdentifier::new("SUM"),
303            distinct: false,
304            args: vec![Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple("col1", false))],
305            order_by: None,
306            filter: None,
307        }];
308
309        let result = extract_aggregates(&exprs, &combined_schema);
310        assert!(result.is_some());
311        let aggregates = result.unwrap();
312        assert_eq!(aggregates.len(), 1);
313        assert!(matches!(aggregates[0].op, AggregateOp::Sum));
314        assert!(matches!(aggregates[0].source, AggregateSource::Column(0)));
315
316        // Test COUNT(*)
317        let exprs = vec![Expression::AggregateFunction {
318            name: vibesql_ast::FunctionIdentifier::new("COUNT"),
319            distinct: false,
320            args: vec![Expression::Wildcard],
321            order_by: None,
322            filter: None,
323        }];
324
325        let result = extract_aggregates(&exprs, &combined_schema);
326        assert!(result.is_some());
327        let aggregates = result.unwrap();
328        assert_eq!(aggregates.len(), 1);
329        assert!(matches!(aggregates[0].op, AggregateOp::Count));
330        assert!(matches!(aggregates[0].source, AggregateSource::CountStar));
331
332        // Test multiple aggregates: SUM(col1), AVG(col2)
333        let exprs = vec![
334            Expression::AggregateFunction {
335                name: vibesql_ast::FunctionIdentifier::new("SUM"),
336                distinct: false,
337                args: vec![Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
338                    "col1", false,
339                ))],
340                order_by: None,
341                filter: None,
342            },
343            Expression::AggregateFunction {
344                name: vibesql_ast::FunctionIdentifier::new("AVG"),
345                distinct: false,
346                args: vec![Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
347                    "col2", false,
348                ))],
349                order_by: None,
350                filter: None,
351            },
352        ];
353
354        let result = extract_aggregates(&exprs, &combined_schema);
355        assert!(result.is_some());
356        let aggregates = result.unwrap();
357        assert_eq!(aggregates.len(), 2);
358        assert!(matches!(aggregates[0].op, AggregateOp::Sum));
359        assert!(matches!(aggregates[0].source, AggregateSource::Column(0)));
360        assert!(matches!(aggregates[1].op, AggregateOp::Avg));
361        assert!(matches!(aggregates[1].source, AggregateSource::Column(1)));
362    }
363
364    #[test]
365    fn test_extract_aggregates_unsupported() {
366        use vibesql_catalog::{ColumnSchema, TableSchema};
367        use vibesql_types::DataType;
368
369        use crate::schema::CombinedSchema;
370
371        let schema = TableSchema::new(
372            "test".to_string(),
373            vec![ColumnSchema::new("col1".to_string(), DataType::Integer, false)],
374        );
375
376        let combined_schema = CombinedSchema::from_table("test".to_string(), schema);
377
378        // Test DISTINCT aggregate (should return None)
379        let exprs = vec![Expression::AggregateFunction {
380            name: vibesql_ast::FunctionIdentifier::new("SUM"),
381            distinct: true,
382            args: vec![Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple("col1", false))],
383            order_by: None,
384            filter: None,
385        }];
386
387        let result = extract_aggregates(&exprs, &combined_schema);
388        assert!(result.is_none());
389
390        // Test non-aggregate expression (should return None)
391        let exprs =
392            vec![Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple("col1", false))];
393
394        let result = extract_aggregates(&exprs, &combined_schema);
395        assert!(result.is_none());
396
397        // Test subquery in aggregate (should return None - not supported)
398        let exprs = vec![Expression::AggregateFunction {
399            name: vibesql_ast::FunctionIdentifier::new("SUM"),
400            distinct: false,
401            args: vec![Expression::ScalarSubquery(Box::new(vibesql_ast::SelectStmt {
402                with_clause: None,
403                distinct: false,
404                select_list: vec![],
405                into_table: None,
406                into_variables: None,
407                from: None,
408                where_clause: None,
409                group_by: None,
410                having: None,
411                window_definitions: None,
412                order_by: None,
413                limit: None,
414                offset: None,
415                set_operation: None,
416                values: None,
417            }))],
418            order_by: None,
419            filter: None,
420        }];
421
422        let result = extract_aggregates(&exprs, &combined_schema);
423        assert!(result.is_none());
424    }
425
426    #[test]
427    fn test_extract_aggregates_with_expression() {
428        use vibesql_catalog::{ColumnSchema, TableSchema};
429        use vibesql_types::DataType;
430
431        use crate::schema::CombinedSchema;
432
433        // Create a simple schema with two columns
434        let schema = TableSchema::new(
435            "test".to_string(),
436            vec![
437                ColumnSchema::new("price".to_string(), DataType::DoublePrecision, false),
438                ColumnSchema::new("discount".to_string(), DataType::DoublePrecision, false),
439            ],
440        );
441
442        let combined_schema = CombinedSchema::from_table("test".to_string(), schema);
443
444        // Test SUM(price * discount) - simple binary operation
445        let exprs = vec![Expression::AggregateFunction {
446            name: vibesql_ast::FunctionIdentifier::new("SUM"),
447            distinct: false,
448            args: vec![Expression::BinaryOp {
449                left: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
450                    "price", false,
451                ))),
452                op: vibesql_ast::BinaryOperator::Multiply,
453                right: Box::new(Expression::ColumnRef(vibesql_ast::ColumnIdentifier::simple(
454                    "discount", false,
455                ))),
456            }],
457            order_by: None,
458            filter: None,
459        }];
460
461        let result = extract_aggregates(&exprs, &combined_schema);
462        assert!(result.is_some());
463        let aggregates = result.unwrap();
464        assert_eq!(aggregates.len(), 1);
465        assert!(matches!(aggregates[0].op, AggregateOp::Sum));
466        assert!(matches!(aggregates[0].source, AggregateSource::Expression(_)));
467    }
468
469    // GROUP BY tests
470
471    #[test]
472    fn test_columnar_group_by_simple() {
473        // Test simple GROUP BY with one group column
474        // SELECT status, SUM(amount) FROM test GROUP BY status
475        let rows = vec![
476            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Double(100.0)]),
477            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B")), SqlValue::Double(200.0)]),
478            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Double(150.0)]),
479            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B")), SqlValue::Double(50.0)]),
480        ];
481
482        let group_cols = vec![0]; // Group by status
483        let agg_cols = vec![(1, AggregateOp::Sum)]; // SUM(amount)
484
485        let result = columnar_group_by(&rows, &group_cols, &agg_cols, None).unwrap();
486
487        // Should have 2 groups: A and B
488        assert_eq!(result.len(), 2);
489
490        // Sort results by group key for deterministic testing
491        let mut sorted = result;
492        sorted.sort_by(|a, b| {
493            let a_key = a.get(0).unwrap();
494            let b_key = b.get(0).unwrap();
495            a_key.partial_cmp(b_key).unwrap()
496        });
497
498        // Check group A: SUM = 250.0
499        assert_eq!(sorted[0].get(0), Some(&SqlValue::Varchar(arcstr::ArcStr::from("A"))));
500        assert!(
501            matches!(sorted[0].get(1), Some(&SqlValue::Double(sum)) if (sum - 250.0).abs() < 0.001)
502        );
503
504        // Check group B: SUM = 250.0
505        assert_eq!(sorted[1].get(0), Some(&SqlValue::Varchar(arcstr::ArcStr::from("B"))));
506        assert!(
507            matches!(sorted[1].get(1), Some(&SqlValue::Double(sum)) if (sum - 250.0).abs() < 0.001)
508        );
509    }
510
511    #[test]
512    fn test_columnar_group_by_multiple_group_keys() {
513        // Test GROUP BY with multiple columns
514        // SELECT status, category, COUNT(*) FROM test GROUP BY status, category
515        let rows = vec![
516            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Integer(1)]),
517            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Integer(2)]),
518            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B")), SqlValue::Integer(1)]),
519            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Integer(1)]),
520            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B")), SqlValue::Integer(2)]),
521        ];
522
523        let group_cols = vec![0, 1]; // Group by status, category
524        let agg_cols = vec![(0, AggregateOp::Count)]; // COUNT(*)
525
526        let result = columnar_group_by(&rows, &group_cols, &agg_cols, None).unwrap();
527
528        // Should have 4 groups: (A,1), (A,2), (B,1), (B,2)
529        assert_eq!(result.len(), 4);
530
531        // Verify we have correct counts
532        for row in &result {
533            let status = row.get(0).unwrap();
534            let category = row.get(1).unwrap();
535            let count = row.get(2).unwrap();
536
537            match (status, category) {
538                (SqlValue::Varchar(s), SqlValue::Integer(1)) if s.as_str() == "A" => {
539                    assert_eq!(count, &SqlValue::Integer(2)); // Two rows with A,1
540                }
541                (SqlValue::Varchar(s), SqlValue::Integer(2)) if s.as_str() == "A" => {
542                    assert_eq!(count, &SqlValue::Integer(1)); // One row with A,2
543                }
544                (SqlValue::Varchar(s), SqlValue::Integer(1)) if s.as_str() == "B" => {
545                    assert_eq!(count, &SqlValue::Integer(1)); // One row with B,1
546                }
547                (SqlValue::Varchar(s), SqlValue::Integer(2)) if s.as_str() == "B" => {
548                    assert_eq!(count, &SqlValue::Integer(1)); // One row with B,2
549                }
550                _ => panic!("Unexpected group key: {:?}, {:?}", status, category),
551            }
552        }
553    }
554
555    #[test]
556    fn test_columnar_group_by_multiple_aggregates() {
557        // Test GROUP BY with multiple aggregate functions
558        // SELECT category, SUM(price), AVG(quantity), COUNT(*) FROM test GROUP BY category
559        let rows = vec![
560            Row::new(vec![SqlValue::Integer(1), SqlValue::Double(100.0), SqlValue::Integer(10)]),
561            Row::new(vec![SqlValue::Integer(2), SqlValue::Double(200.0), SqlValue::Integer(20)]),
562            Row::new(vec![SqlValue::Integer(1), SqlValue::Double(150.0), SqlValue::Integer(15)]),
563        ];
564
565        let group_cols = vec![0]; // Group by category
566        let agg_cols = vec![
567            (1, AggregateOp::Sum),   // SUM(price)
568            (2, AggregateOp::Avg),   // AVG(quantity)
569            (0, AggregateOp::Count), // COUNT(*)
570        ];
571
572        let result = columnar_group_by(&rows, &group_cols, &agg_cols, None).unwrap();
573
574        // Should have 2 groups
575        assert_eq!(result.len(), 2);
576
577        // Sort by category for deterministic testing
578        let mut sorted = result;
579        sorted.sort_by(|a, b| {
580            let a_key = a.get(0).unwrap();
581            let b_key = b.get(0).unwrap();
582            a_key.partial_cmp(b_key).unwrap()
583        });
584
585        // Group 1: SUM=250, AVG=12.5, COUNT=2
586        assert_eq!(sorted[0].get(0), Some(&SqlValue::Integer(1)));
587        assert!(
588            matches!(sorted[0].get(1), Some(&SqlValue::Double(sum)) if (sum - 250.0).abs() < 0.001)
589        );
590        assert!(
591            matches!(sorted[0].get(2), Some(&SqlValue::Double(avg)) if (avg - 12.5).abs() < 0.001)
592        );
593        assert_eq!(sorted[0].get(3), Some(&SqlValue::Integer(2)));
594
595        // Group 2: SUM=200, AVG=20.0, COUNT=1
596        assert_eq!(sorted[1].get(0), Some(&SqlValue::Integer(2)));
597        assert!(
598            matches!(sorted[1].get(1), Some(&SqlValue::Double(sum)) if (sum - 200.0).abs() < 0.001)
599        );
600        assert!(
601            matches!(sorted[1].get(2), Some(&SqlValue::Double(avg)) if (avg - 20.0).abs() < 0.001)
602        );
603        assert_eq!(sorted[1].get(3), Some(&SqlValue::Integer(1)));
604    }
605
606    #[test]
607    fn test_columnar_group_by_with_filter() {
608        // Test GROUP BY with pre-filtering
609        // SELECT status, SUM(amount) FROM test WHERE amount > 100 GROUP BY status
610        let rows = vec![
611            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Double(100.0)]),
612            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B")), SqlValue::Double(200.0)]),
613            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Double(150.0)]),
614            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("B")), SqlValue::Double(50.0)]),
615        ];
616
617        // Filter: amount > 100 (rows 1 and 2)
618        let filter = vec![false, true, true, false];
619
620        let group_cols = vec![0]; // Group by status
621        let agg_cols = vec![(1, AggregateOp::Sum)]; // SUM(amount)
622
623        let result = columnar_group_by(&rows, &group_cols, &agg_cols, Some(&filter)).unwrap();
624
625        // Should have 2 groups (only rows passing filter)
626        assert_eq!(result.len(), 2);
627
628        // Sort results by group key
629        let mut sorted = result;
630        sorted.sort_by(|a, b| {
631            let a_key = a.get(0).unwrap();
632            let b_key = b.get(0).unwrap();
633            a_key.partial_cmp(b_key).unwrap()
634        });
635
636        // Check group A: only row 2 (150.0) passes filter
637        assert_eq!(sorted[0].get(0), Some(&SqlValue::Varchar(arcstr::ArcStr::from("A"))));
638        assert!(
639            matches!(sorted[0].get(1), Some(&SqlValue::Double(sum)) if (sum - 150.0).abs() < 0.001)
640        );
641
642        // Check group B: only row 1 (200.0) passes filter
643        assert_eq!(sorted[1].get(0), Some(&SqlValue::Varchar(arcstr::ArcStr::from("B"))));
644        assert!(
645            matches!(sorted[1].get(1), Some(&SqlValue::Double(sum)) if (sum - 200.0).abs() < 0.001)
646        );
647    }
648
649    #[test]
650    fn test_columnar_group_by_empty_input() {
651        let rows: Vec<Row> = vec![];
652        let group_cols = vec![0];
653        let agg_cols = vec![(1, AggregateOp::Sum)];
654
655        let result = columnar_group_by(&rows, &group_cols, &agg_cols, None).unwrap();
656
657        // Should return empty result for empty input
658        assert_eq!(result.len(), 0);
659    }
660
661    #[test]
662    fn test_columnar_group_by_null_in_group_key() {
663        // Test that NULL values in group keys are handled correctly
664        let rows = vec![
665            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Double(100.0)]),
666            Row::new(vec![SqlValue::Null, SqlValue::Double(200.0)]),
667            Row::new(vec![SqlValue::Varchar(arcstr::ArcStr::from("A")), SqlValue::Double(150.0)]),
668            Row::new(vec![SqlValue::Null, SqlValue::Double(50.0)]),
669        ];
670
671        let group_cols = vec![0]; // Group by first column
672        let agg_cols = vec![(1, AggregateOp::Sum)]; // SUM
673
674        let result = columnar_group_by(&rows, &group_cols, &agg_cols, None).unwrap();
675
676        // Should have 2 groups: "A" and NULL
677        assert_eq!(result.len(), 2);
678
679        // Find the groups
680        let a_group = result
681            .iter()
682            .find(|r| matches!(r.get(0), Some(SqlValue::Varchar(s)) if s.as_str() == "A"));
683        let null_group = result.iter().find(|r| matches!(r.get(0), Some(SqlValue::Null)));
684
685        assert!(a_group.is_some());
686        assert!(null_group.is_some());
687
688        // Check "A" group: 100 + 150 = 250
689        assert!(
690            matches!(a_group.unwrap().get(1), Some(&SqlValue::Double(sum)) if (sum - 250.0).abs() < 0.001)
691        );
692
693        // Check NULL group: 200 + 50 = 250
694        assert!(
695            matches!(null_group.unwrap().get(1), Some(&SqlValue::Double(sum)) if (sum - 250.0).abs() < 0.001)
696        );
697    }
698}
vibesql_executor/select/columnar/aggregate/mod.rs

vibesql_executor/select/columnar/aggregate/
mod.rs
