alopex_sql/executor/query/
aggregate.rs

1use std::cmp::Ordering;
2use std::collections::{HashMap, HashSet};
3
4use crate::catalog::ColumnMetadata;
5use crate::executor::evaluator::EvalContext;
6use crate::executor::memory::{MemoryPolicy, MemoryTracker, map_core_memory_error};
7use crate::executor::{ExecutorError, Result};
8use crate::planner::aggregate_expr::{AggregateExpr, AggregateFunction};
9use crate::planner::typed_expr::TypedExpr;
10use crate::storage::{RowCodec, SqlValue};
11
12use super::{Row, RowIterator};
13
14/// Byte-encoded group key for hash-based aggregation.
15pub type GroupKeyBytes = Vec<u8>;
16
17fn encode_group_value(value: &SqlValue, buf: &mut Vec<u8>) -> Result<()> {
18    buf.push(value.type_tag());
19    match value {
20        SqlValue::Null => Ok(()),
21        SqlValue::Integer(v) => {
22            buf.extend_from_slice(&v.to_le_bytes());
23            Ok(())
24        }
25        SqlValue::BigInt(v) => {
26            buf.extend_from_slice(&v.to_le_bytes());
27            Ok(())
28        }
29        SqlValue::Float(v) => {
30            buf.extend_from_slice(&v.to_bits().to_le_bytes());
31            Ok(())
32        }
33        SqlValue::Double(v) => {
34            buf.extend_from_slice(&v.to_bits().to_le_bytes());
35            Ok(())
36        }
37        SqlValue::Text(s) => {
38            let len = u32::try_from(s.len()).map_err(|_| ExecutorError::InvalidOperation {
39                operation: "aggregate".into(),
40                reason: "text length exceeds u32::MAX".into(),
41            })?;
42            buf.extend_from_slice(&len.to_le_bytes());
43            buf.extend_from_slice(s.as_bytes());
44            Ok(())
45        }
46        SqlValue::Blob(bytes) => {
47            let len = u32::try_from(bytes.len()).map_err(|_| ExecutorError::InvalidOperation {
48                operation: "aggregate".into(),
49                reason: "blob length exceeds u32::MAX".into(),
50            })?;
51            buf.extend_from_slice(&len.to_le_bytes());
52            buf.extend_from_slice(bytes);
53            Ok(())
54        }
55        SqlValue::Boolean(b) => {
56            buf.push(u8::from(*b));
57            Ok(())
58        }
59        SqlValue::Timestamp(v) => {
60            buf.extend_from_slice(&v.to_le_bytes());
61            Ok(())
62        }
63        SqlValue::Vector(values) => {
64            let len = u32::try_from(values.len()).map_err(|_| ExecutorError::InvalidOperation {
65                operation: "aggregate".into(),
66                reason: "vector length exceeds u32::MAX".into(),
67            })?;
68            buf.extend_from_slice(&len.to_le_bytes());
69            for f in values {
70                buf.extend_from_slice(&f.to_bits().to_le_bytes());
71            }
72            Ok(())
73        }
74    }
75}
76
77/// Encode group key values into a deterministic byte sequence.
78pub fn encode_group_key(values: &[SqlValue]) -> Result<GroupKeyBytes> {
79    let mut buf = Vec::new();
80    for value in values {
81        encode_group_value(value, &mut buf)?;
82    }
83    Ok(buf)
84}
85
86/// Accumulator interface for aggregate function execution.
87pub trait Accumulator: Send {
88    /// Update the accumulator with a new value (None for COUNT(*) rows).
89    fn update(&mut self, value: Option<SqlValue>) -> Result<()>;
90    /// Finalize the accumulator and return the resulting SqlValue.
91    fn finalize(&self) -> Result<SqlValue>;
92    /// Clone the accumulator as a trait object.
93    fn clone_box(&self) -> Box<dyn Accumulator>;
94}
95
96impl Clone for Box<dyn Accumulator> {
97    fn clone(&self) -> Self {
98        self.clone_box()
99    }
100}
101
102/// Accumulator for COUNT / COUNT(DISTINCT).
103#[derive(Debug, Clone)]
104pub struct CountAccumulator {
105    count: usize,
106    distinct_values: Option<HashSet<Vec<u8>>>,
107}
108
109impl CountAccumulator {
110    /// Create a new count accumulator.
111    pub fn new(distinct: bool) -> Self {
112        Self {
113            count: 0,
114            distinct_values: if distinct { Some(HashSet::new()) } else { None },
115        }
116    }
117}
118
119impl Accumulator for CountAccumulator {
120    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
121        match (&mut self.distinct_values, value) {
122            (Some(distinct), Some(value)) => {
123                if value.is_null() {
124                    return Ok(());
125                }
126                let encoded = RowCodec::encode(std::slice::from_ref(&value));
127                if distinct.insert(encoded) {
128                    self.count += 1;
129                }
130            }
131            (Some(_), None) => {
132                self.count += 1;
133            }
134            (None, Some(value)) => {
135                if !value.is_null() {
136                    self.count += 1;
137                }
138            }
139            (None, None) => {
140                self.count += 1;
141            }
142        }
143        Ok(())
144    }
145
146    fn finalize(&self) -> Result<SqlValue> {
147        Ok(SqlValue::BigInt(self.count as i64))
148    }
149
150    fn clone_box(&self) -> Box<dyn Accumulator> {
151        Box::new(self.clone())
152    }
153}
154
155/// Accumulator for SUM.
156#[derive(Debug, Clone)]
157pub struct SumAccumulator {
158    sum: Option<f64>,
159}
160
161impl SumAccumulator {
162    /// Create a new sum accumulator.
163    pub fn new() -> Self {
164        Self { sum: None }
165    }
166}
167
168impl Default for SumAccumulator {
169    fn default() -> Self {
170        Self::new()
171    }
172}
173
174impl Accumulator for SumAccumulator {
175    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
176        let Some(value) = value else {
177            return Ok(());
178        };
179        if value.is_null() {
180            return Ok(());
181        }
182        let numeric = numeric_to_f64(&value)?;
183        self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
184        Ok(())
185    }
186
187    fn finalize(&self) -> Result<SqlValue> {
188        Ok(self.sum.map_or(SqlValue::Null, SqlValue::Double))
189    }
190
191    fn clone_box(&self) -> Box<dyn Accumulator> {
192        Box::new(self.clone())
193    }
194}
195
196/// Accumulator for TOTAL (SUM that returns 0.0 on empty/all-NULL input).
197#[derive(Debug, Clone)]
198pub struct TotalAccumulator {
199    sum: Option<f64>,
200}
201
202impl TotalAccumulator {
203    /// Create a new total accumulator.
204    pub fn new() -> Self {
205        Self { sum: None }
206    }
207}
208
209impl Default for TotalAccumulator {
210    fn default() -> Self {
211        Self::new()
212    }
213}
214
215impl Accumulator for TotalAccumulator {
216    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
217        let Some(value) = value else {
218            return Ok(());
219        };
220        if value.is_null() {
221            return Ok(());
222        }
223        let numeric = numeric_to_f64(&value)?;
224        self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
225        Ok(())
226    }
227
228    fn finalize(&self) -> Result<SqlValue> {
229        Ok(SqlValue::Double(self.sum.unwrap_or(0.0)))
230    }
231
232    fn clone_box(&self) -> Box<dyn Accumulator> {
233        Box::new(self.clone())
234    }
235}
236
237/// Accumulator for AVG.
238#[derive(Debug, Clone)]
239pub struct AvgAccumulator {
240    sum: Option<f64>,
241    count: usize,
242}
243
244impl AvgAccumulator {
245    /// Create a new average accumulator.
246    pub fn new() -> Self {
247        Self {
248            sum: None,
249            count: 0,
250        }
251    }
252}
253
254impl Default for AvgAccumulator {
255    fn default() -> Self {
256        Self::new()
257    }
258}
259
260impl Accumulator for AvgAccumulator {
261    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
262        let Some(value) = value else {
263            return Ok(());
264        };
265        if value.is_null() {
266            return Ok(());
267        }
268        let numeric = numeric_to_f64(&value)?;
269        self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
270        self.count += 1;
271        Ok(())
272    }
273
274    fn finalize(&self) -> Result<SqlValue> {
275        if self.count == 0 {
276            return Ok(SqlValue::Null);
277        }
278        let sum = self.sum.unwrap_or(0.0);
279        Ok(SqlValue::Double(sum / self.count as f64))
280    }
281
282    fn clone_box(&self) -> Box<dyn Accumulator> {
283        Box::new(self.clone())
284    }
285}
286
287fn numeric_to_f64(value: &SqlValue) -> Result<f64> {
288    match value {
289        SqlValue::Integer(v) => Ok(*v as f64),
290        SqlValue::BigInt(v) => Ok(*v as f64),
291        SqlValue::Float(v) => Ok(*v as f64),
292        SqlValue::Double(v) => Ok(*v),
293        _ => Err(ExecutorError::Evaluation(
294            crate::executor::EvaluationError::TypeMismatch {
295                expected: "numeric".into(),
296                actual: value.type_name().into(),
297            },
298        )),
299    }
300}
301
302/// Accumulator for MIN / MAX.
303#[derive(Debug, Clone)]
304pub struct MinMaxAccumulator {
305    value: Option<SqlValue>,
306    is_min: bool,
307}
308
309impl MinMaxAccumulator {
310    /// Create a new min/max accumulator.
311    pub fn new(is_min: bool) -> Self {
312        Self {
313            value: None,
314            is_min,
315        }
316    }
317}
318
319impl Accumulator for MinMaxAccumulator {
320    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
321        let Some(value) = value else {
322            return Ok(());
323        };
324        if value.is_null() {
325            return Ok(());
326        }
327
328        match &self.value {
329            None => {
330                self.value = Some(value);
331            }
332            Some(current) => {
333                if std::mem::discriminant(current) != std::mem::discriminant(&value) {
334                    return Err(ExecutorError::Evaluation(
335                        crate::executor::EvaluationError::TypeMismatch {
336                            expected: current.type_name().into(),
337                            actual: value.type_name().into(),
338                        },
339                    ));
340                }
341                let ordering = value.partial_cmp(current).ok_or_else(|| {
342                    ExecutorError::Evaluation(crate::executor::EvaluationError::TypeMismatch {
343                        expected: current.type_name().into(),
344                        actual: value.type_name().into(),
345                    })
346                })?;
347                let should_replace = matches!(
348                    (self.is_min, ordering),
349                    (true, Ordering::Less) | (false, Ordering::Greater)
350                );
351                if should_replace {
352                    self.value = Some(value);
353                }
354            }
355        }
356        Ok(())
357    }
358
359    fn finalize(&self) -> Result<SqlValue> {
360        Ok(self.value.clone().unwrap_or(SqlValue::Null))
361    }
362
363    fn clone_box(&self) -> Box<dyn Accumulator> {
364        Box::new(self.clone())
365    }
366}
367
368/// Accumulator for GROUP_CONCAT.
369#[derive(Debug, Clone)]
370pub struct GroupConcatAccumulator {
371    values: Vec<String>,
372    separator: String,
373}
374
375impl GroupConcatAccumulator {
376    /// Create a new GROUP_CONCAT accumulator with the given separator.
377    pub fn new(separator: String) -> Self {
378        Self {
379            values: Vec::new(),
380            separator,
381        }
382    }
383}
384
385impl Accumulator for GroupConcatAccumulator {
386    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
387        let Some(value) = value else {
388            return Ok(());
389        };
390        match value {
391            SqlValue::Null => Ok(()),
392            SqlValue::Text(text) => {
393                self.values.push(text);
394                Ok(())
395            }
396            other => Err(ExecutorError::Evaluation(
397                crate::executor::EvaluationError::TypeMismatch {
398                    expected: "Text".into(),
399                    actual: other.type_name().into(),
400                },
401            )),
402        }
403    }
404
405    fn finalize(&self) -> Result<SqlValue> {
406        if self.values.is_empty() {
407            return Ok(SqlValue::Null);
408        }
409        Ok(SqlValue::Text(self.values.join(&self.separator)))
410    }
411
412    fn clone_box(&self) -> Box<dyn Accumulator> {
413        Box::new(self.clone())
414    }
415}
416
417/// Accumulator for STRING_AGG.
418#[derive(Debug, Clone)]
419pub struct StringAggAccumulator {
420    values: Vec<String>,
421    separator: String,
422}
423
424impl StringAggAccumulator {
425    /// Create a new string_agg accumulator.
426    pub fn new(separator: String) -> Self {
427        Self {
428            values: Vec::new(),
429            separator,
430        }
431    }
432}
433
434impl Accumulator for StringAggAccumulator {
435    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
436        let Some(value) = value else {
437            return Ok(());
438        };
439        match value {
440            SqlValue::Null => Ok(()),
441            SqlValue::Text(s) => {
442                self.values.push(s);
443                Ok(())
444            }
445            other => Err(ExecutorError::Evaluation(
446                crate::executor::EvaluationError::TypeMismatch {
447                    expected: "Text".into(),
448                    actual: other.type_name().into(),
449                },
450            )),
451        }
452    }
453
454    fn finalize(&self) -> Result<SqlValue> {
455        if self.values.is_empty() {
456            return Ok(SqlValue::Null);
457        }
458        Ok(SqlValue::Text(self.values.join(&self.separator)))
459    }
460
461    fn clone_box(&self) -> Box<dyn Accumulator> {
462        Box::new(self.clone())
463    }
464}
465
466/// Create a new accumulator instance for the aggregate function.
467pub fn create_accumulator(function: &AggregateFunction, distinct: bool) -> Box<dyn Accumulator> {
468    match function {
469        AggregateFunction::Count => Box::new(CountAccumulator::new(distinct)),
470        AggregateFunction::Sum => Box::new(SumAccumulator::new()),
471        AggregateFunction::Total => Box::new(TotalAccumulator::new()),
472        AggregateFunction::Avg => Box::new(AvgAccumulator::new()),
473        AggregateFunction::Min => Box::new(MinMaxAccumulator::new(true)),
474        AggregateFunction::Max => Box::new(MinMaxAccumulator::new(false)),
475        AggregateFunction::GroupConcat { separator } => {
476            let sep = separator.clone().unwrap_or_else(|| ",".to_string());
477            Box::new(GroupConcatAccumulator::new(sep))
478        }
479        AggregateFunction::StringAgg { separator } => {
480            let sep = separator.clone().unwrap_or_else(|| ",".to_string());
481            Box::new(StringAggAccumulator::new(sep))
482        }
483    }
484}
485
486const DEFAULT_GROUP_LIMIT: usize = 1_000_000;
487const AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES: u64 = 32;
488
489struct AggregateGroup {
490    key_values: Vec<SqlValue>,
491    accumulators: Vec<Box<dyn Accumulator>>,
492}
493
494/// Iterator that performs hash-based aggregation over input rows.
495pub struct AggregateIterator<'a> {
496    input: Box<dyn RowIterator + 'a>,
497    group_keys: Vec<TypedExpr>,
498    aggregates: Vec<AggregateExpr>,
499    having: Option<TypedExpr>,
500    hash_table: Option<HashMap<GroupKeyBytes, AggregateGroup>>,
501    result_rows: Vec<Row>,
502    index: usize,
503    schema: Vec<ColumnMetadata>,
504    group_limit: usize,
505    memory_tracker: Option<MemoryTracker>,
506}
507
508impl<'a> AggregateIterator<'a> {
509    /// Create a new aggregate iterator with the default group limit.
510    pub fn new(
511        input: Box<dyn RowIterator + 'a>,
512        group_keys: Vec<TypedExpr>,
513        aggregates: Vec<AggregateExpr>,
514        having: Option<TypedExpr>,
515        schema: Vec<ColumnMetadata>,
516    ) -> Self {
517        Self {
518            input,
519            group_keys,
520            aggregates,
521            having,
522            hash_table: None,
523            result_rows: Vec::new(),
524            index: 0,
525            schema,
526            group_limit: DEFAULT_GROUP_LIMIT,
527            memory_tracker: None,
528        }
529    }
530
531    /// Override the maximum number of groups allowed during aggregation.
532    pub fn with_group_limit(mut self, limit: usize) -> Self {
533        self.group_limit = limit;
534        self
535    }
536
537    /// Attach a memory policy for enforcing in-flight aggregation limits.
538    pub fn with_memory_policy(mut self, policy: Option<MemoryPolicy>) -> Self {
539        self.memory_tracker = policy.map(MemoryTracker::new);
540        self
541    }
542
543    fn build_hash_table(&mut self) -> Result<()> {
544        let mut table: HashMap<GroupKeyBytes, AggregateGroup> = HashMap::new();
545        let mut next_row_id = 0u64;
546
547        while let Some(result) = self.input.next_row() {
548            let row = result?;
549            let ctx = EvalContext::new(&row.values);
550
551            let mut key_values = Vec::with_capacity(self.group_keys.len());
552            for expr in &self.group_keys {
553                key_values.push(crate::executor::evaluator::evaluate(expr, &ctx)?);
554            }
555            let key_bytes = encode_group_key(&key_values)?;
556
557            if !table.contains_key(&key_bytes) {
558                if table.len() + 1 > self.group_limit {
559                    return Err(ExecutorError::ResourceExhausted {
560                        message: format!(
561                            "GROUP BY result exceeds memory limit (max groups: {})",
562                            self.group_limit
563                        ),
564                    });
565                }
566                if let Some(tracker) = &mut self.memory_tracker {
567                    tracker
568                        .add_values(&key_values)
569                        .map_err(map_core_memory_error)?;
570                    tracker
571                        .add_bytes(
572                            self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES,
573                        )
574                        .map_err(map_core_memory_error)?;
575                }
576                let accumulators = self
577                    .aggregates
578                    .iter()
579                    .map(|agg| create_accumulator(&agg.function, agg.distinct))
580                    .collect::<Vec<_>>();
581                table.insert(
582                    key_bytes.clone(),
583                    AggregateGroup {
584                        key_values: key_values.clone(),
585                        accumulators,
586                    },
587                );
588            }
589
590            if let Some(group) = table.get_mut(&key_bytes) {
591                for (idx, agg) in self.aggregates.iter().enumerate() {
592                    let value = match &agg.arg {
593                        None => None,
594                        Some(expr) => Some(crate::executor::evaluator::evaluate(expr, &ctx)?),
595                    };
596                    if let Some(tracker) = &mut self.memory_tracker
597                        && matches!(
598                            agg.function,
599                            AggregateFunction::GroupConcat { .. }
600                                | AggregateFunction::StringAgg { .. }
601                        )
602                        && let Some(value_ref) = value.as_ref()
603                    {
604                        tracker
605                            .add_value(value_ref)
606                            .map_err(map_core_memory_error)?;
607                    }
608                    group.accumulators[idx].update(value)?;
609                }
610            }
611        }
612
613        if table.is_empty() && self.group_keys.is_empty() {
614            if let Some(tracker) = &mut self.memory_tracker {
615                tracker
616                    .add_bytes(self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES)
617                    .map_err(map_core_memory_error)?;
618            }
619            let accumulators = self
620                .aggregates
621                .iter()
622                .map(|agg| create_accumulator(&agg.function, agg.distinct))
623                .collect::<Vec<_>>();
624            table.insert(
625                Vec::new(),
626                AggregateGroup {
627                    key_values: Vec::new(),
628                    accumulators,
629                },
630            );
631        }
632
633        let mut rows = Vec::with_capacity(table.len());
634        for group in table.values() {
635            let mut values = Vec::with_capacity(self.group_keys.len() + self.aggregates.len());
636            values.extend(group.key_values.iter().cloned());
637            for acc in &group.accumulators {
638                values.push(acc.finalize()?);
639            }
640            let row = Row::new(next_row_id, values);
641            next_row_id += 1;
642            if let Some(tracker) = &mut self.memory_tracker {
643                tracker
644                    .add_row(&row.values)
645                    .map_err(map_core_memory_error)?;
646            }
647
648            if let Some(having) = &self.having {
649                let ctx = EvalContext::new(&row.values);
650                match crate::executor::evaluator::evaluate(having, &ctx)? {
651                    SqlValue::Boolean(true) => rows.push(row),
652                    SqlValue::Boolean(false) | SqlValue::Null => {}
653                    other => {
654                        return Err(ExecutorError::Evaluation(
655                            crate::executor::EvaluationError::TypeMismatch {
656                                expected: "Boolean".into(),
657                                actual: other.type_name().into(),
658                            },
659                        ));
660                    }
661                }
662            } else {
663                rows.push(row);
664            }
665        }
666
667        self.hash_table = Some(table);
668        self.result_rows = rows;
669        Ok(())
670    }
671}
672
673impl<'a> RowIterator for AggregateIterator<'a> {
674    fn next_row(&mut self) -> Option<Result<Row>> {
675        if self.hash_table.is_none()
676            && let Err(err) = self.build_hash_table()
677        {
678            return Some(Err(err));
679        }
680
681        if self.index >= self.result_rows.len() {
682            return None;
683        }
684        let row = self.result_rows[self.index].clone();
685        self.index += 1;
686        Some(Ok(row))
687    }
688
689    fn schema(&self) -> &[ColumnMetadata] {
690        &self.schema
691    }
692}
693
694/// Iterator that performs streaming aggregation over sorted input.
695pub struct StreamingAggregateIterator<'a> {
696    input: Box<dyn RowIterator + 'a>,
697    group_keys: Vec<TypedExpr>,
698    aggregates: Vec<AggregateExpr>,
699    having: Option<TypedExpr>,
700    schema: Vec<ColumnMetadata>,
701    current_key: Option<Vec<SqlValue>>,
702    accumulators: Vec<Box<dyn Accumulator>>,
703    pending_row: Option<Row>,
704    finished: bool,
705    next_row_id: u64,
706    saw_row: bool,
707}
708
709impl<'a> StreamingAggregateIterator<'a> {
710    pub fn new(
711        input: Box<dyn RowIterator + 'a>,
712        group_keys: Vec<TypedExpr>,
713        aggregates: Vec<AggregateExpr>,
714        having: Option<TypedExpr>,
715        schema: Vec<ColumnMetadata>,
716    ) -> Self {
717        Self {
718            input,
719            group_keys,
720            aggregates,
721            having,
722            schema,
723            current_key: None,
724            accumulators: Vec::new(),
725            pending_row: None,
726            finished: false,
727            next_row_id: 0,
728            saw_row: false,
729        }
730    }
731
732    fn init_accumulators(&self) -> Vec<Box<dyn Accumulator>> {
733        self.aggregates
734            .iter()
735            .map(|agg| create_accumulator(&agg.function, agg.distinct))
736            .collect()
737    }
738
739    fn update_accumulators(&mut self, ctx: &EvalContext<'_>) -> Result<()> {
740        for (idx, agg) in self.aggregates.iter().enumerate() {
741            let value = match &agg.arg {
742                None => None,
743                Some(expr) => Some(crate::executor::evaluator::evaluate(expr, ctx)?),
744            };
745            self.accumulators[idx].update(value)?;
746        }
747        Ok(())
748    }
749
750    fn finalize_group(&mut self, key_values: &[SqlValue]) -> Result<Option<Row>> {
751        let mut values = Vec::with_capacity(self.group_keys.len() + self.aggregates.len());
752        values.extend(key_values.iter().cloned());
753        for acc in &self.accumulators {
754            values.push(acc.finalize()?);
755        }
756        let row = Row::new(self.next_row_id, values);
757        self.next_row_id = self.next_row_id.saturating_add(1);
758
759        if let Some(having) = &self.having {
760            let ctx = EvalContext::new(&row.values);
761            match crate::executor::evaluator::evaluate(having, &ctx)? {
762                SqlValue::Boolean(true) => Ok(Some(row)),
763                SqlValue::Boolean(false) | SqlValue::Null => Ok(None),
764                other => Err(ExecutorError::Evaluation(
765                    crate::executor::EvaluationError::TypeMismatch {
766                        expected: "Boolean".into(),
767                        actual: other.type_name().into(),
768                    },
769                )),
770            }
771        } else {
772            Ok(Some(row))
773        }
774    }
775}
776
777impl<'a> RowIterator for StreamingAggregateIterator<'a> {
778    fn next_row(&mut self) -> Option<Result<Row>> {
779        if let Some(row) = self.pending_row.take() {
780            return Some(Ok(row));
781        }
782        if self.finished {
783            return None;
784        }
785
786        loop {
787            match self.input.next_row() {
788                Some(Ok(row)) => {
789                    self.saw_row = true;
790                    let ctx = EvalContext::new(&row.values);
791                    let mut key_values = Vec::with_capacity(self.group_keys.len());
792                    for expr in &self.group_keys {
793                        match crate::executor::evaluator::evaluate(expr, &ctx) {
794                            Ok(value) => key_values.push(value),
795                            Err(err) => return Some(Err(err)),
796                        }
797                    }
798
799                    match &self.current_key {
800                        None => {
801                            self.current_key = Some(key_values);
802                            self.accumulators = self.init_accumulators();
803                            if let Err(err) = self.update_accumulators(&ctx) {
804                                return Some(Err(err));
805                            }
806                        }
807                        Some(current_key) if *current_key == key_values => {
808                            if let Err(err) = self.update_accumulators(&ctx) {
809                                return Some(Err(err));
810                            }
811                        }
812                        Some(_) => {
813                            let current_key = self.current_key.clone().unwrap_or_default();
814                            let output = match self.finalize_group(&current_key) {
815                                Ok(value) => value,
816                                Err(err) => return Some(Err(err)),
817                            };
818                            self.current_key = Some(key_values);
819                            self.accumulators = self.init_accumulators();
820                            if let Err(err) = self.update_accumulators(&ctx) {
821                                return Some(Err(err));
822                            }
823                            if let Some(row) = output {
824                                return Some(Ok(row));
825                            }
826                        }
827                    }
828                }
829                Some(Err(err)) => return Some(Err(err)),
830                None => {
831                    self.finished = true;
832                    if let Some(current_key) = self.current_key.take() {
833                        return match self.finalize_group(&current_key) {
834                            Ok(Some(row)) => Some(Ok(row)),
835                            Ok(None) => None,
836                            Err(err) => Some(Err(err)),
837                        };
838                    }
839
840                    if self.group_keys.is_empty() && !self.saw_row {
841                        self.accumulators = self.init_accumulators();
842                        return match self.finalize_group(&[]) {
843                            Ok(Some(row)) => Some(Ok(row)),
844                            Ok(None) => None,
845                            Err(err) => Some(Err(err)),
846                        };
847                    }
848
849                    return None;
850                }
851            }
852        }
853    }
854
855    fn schema(&self) -> &[ColumnMetadata] {
856        &self.schema
857    }
858}
859
860/// Build output schema for aggregate results.
861pub fn build_aggregate_schema(
862    group_keys: &[TypedExpr],
863    aggregates: &[AggregateExpr],
864) -> Vec<ColumnMetadata> {
865    let mut schema = Vec::new();
866    for (idx, key) in group_keys.iter().enumerate() {
867        let name = match &key.kind {
868            crate::planner::typed_expr::TypedExprKind::ColumnRef { column, .. } => column.clone(),
869            _ => format!("group_{idx}"),
870        };
871        schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
872    }
873    for (idx, agg) in aggregates.iter().enumerate() {
874        let name = match &agg.function {
875            AggregateFunction::Count => format!("count_{idx}"),
876            AggregateFunction::Sum => format!("sum_{idx}"),
877            AggregateFunction::Total => format!("total_{idx}"),
878            AggregateFunction::Avg => format!("avg_{idx}"),
879            AggregateFunction::Min => format!("min_{idx}"),
880            AggregateFunction::Max => format!("max_{idx}"),
881            AggregateFunction::GroupConcat { .. } => format!("group_concat_{idx}"),
882            AggregateFunction::StringAgg { .. } => format!("string_agg_{idx}"),
883        };
884        schema.push(ColumnMetadata::new(name, agg.result_type.clone()));
885    }
886    schema
887}
888
889#[cfg(test)]
890mod tests {
891    use super::*;
892
893    #[test]
894    fn count_accumulator_counts_rows_and_skips_nulls() {
895        let mut acc = CountAccumulator::new(false);
896        acc.update(None).unwrap();
897        acc.update(Some(SqlValue::Null)).unwrap();
898        acc.update(Some(SqlValue::Integer(1))).unwrap();
899        assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(2));
900    }
901
902    #[test]
903    fn count_accumulator_distinct_deduplicates() {
904        let mut acc = CountAccumulator::new(true);
905        acc.update(Some(SqlValue::Integer(1))).unwrap();
906        acc.update(Some(SqlValue::Integer(1))).unwrap();
907        acc.update(Some(SqlValue::Integer(2))).unwrap();
908        assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(2));
909    }
910
911    #[test]
912    fn sum_accumulator_aggregates_numeric_values() {
913        let mut acc = SumAccumulator::new();
914        acc.update(Some(SqlValue::Integer(2))).unwrap();
915        acc.update(Some(SqlValue::Double(3.5))).unwrap();
916        acc.update(Some(SqlValue::Null)).unwrap();
917        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(5.5));
918    }
919
920    #[test]
921    fn total_accumulator_returns_zero_for_empty() {
922        let acc = TotalAccumulator::new();
923        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(0.0));
924    }
925
926    #[test]
927    fn total_accumulator_aggregates_numeric_values() {
928        let mut acc = TotalAccumulator::new();
929        acc.update(Some(SqlValue::Integer(2))).unwrap();
930        acc.update(Some(SqlValue::Null)).unwrap();
931        acc.update(Some(SqlValue::Double(1.5))).unwrap();
932        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(3.5));
933    }
934
935    #[test]
936    fn avg_accumulator_handles_empty_and_nulls() {
937        let mut acc = AvgAccumulator::new();
938        assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
939        acc.update(Some(SqlValue::Null)).unwrap();
940        acc.update(Some(SqlValue::BigInt(4))).unwrap();
941        acc.update(Some(SqlValue::Integer(2))).unwrap();
942        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(3.0));
943    }
944
945    #[test]
946    fn min_max_accumulator_tracks_extremes() {
947        let mut min_acc = MinMaxAccumulator::new(true);
948        let mut max_acc = MinMaxAccumulator::new(false);
949        for value in [3, 1, 2] {
950            min_acc.update(Some(SqlValue::Integer(value))).unwrap();
951            max_acc.update(Some(SqlValue::Integer(value))).unwrap();
952        }
953        assert_eq!(min_acc.finalize().unwrap(), SqlValue::Integer(1));
954        assert_eq!(max_acc.finalize().unwrap(), SqlValue::Integer(3));
955    }
956
957    #[test]
958    fn min_max_accumulator_rejects_type_mismatch() {
959        let mut acc = MinMaxAccumulator::new(true);
960        acc.update(Some(SqlValue::Integer(1))).unwrap();
961        let err = acc.update(Some(SqlValue::Text("bad".into()))).unwrap_err();
962        match err {
963            ExecutorError::Evaluation(crate::executor::EvaluationError::TypeMismatch {
964                ..
965            }) => {}
966            other => panic!("unexpected error {:?}", other),
967        }
968    }
969
970    #[test]
971    fn group_concat_accumulator_joins_values() {
972        let mut acc = GroupConcatAccumulator::new("|".into());
973        acc.update(Some(SqlValue::Text("a".into()))).unwrap();
974        acc.update(Some(SqlValue::Null)).unwrap();
975        acc.update(Some(SqlValue::Text("b".into()))).unwrap();
976        assert_eq!(acc.finalize().unwrap(), SqlValue::Text("a|b".into()));
977    }
978
979    #[test]
980    fn group_concat_accumulator_empty_returns_null() {
981        let acc = GroupConcatAccumulator::new(",".into());
982        assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
983    }
984
985    #[test]
986    fn string_agg_accumulator_joins_values() {
987        let mut acc = StringAggAccumulator::new("::".into());
988        acc.update(Some(SqlValue::Text("a".into()))).unwrap();
989        acc.update(Some(SqlValue::Null)).unwrap();
990        acc.update(Some(SqlValue::Text("b".into()))).unwrap();
991        assert_eq!(acc.finalize().unwrap(), SqlValue::Text("a::b".into()));
992    }
993
994    #[test]
995    fn string_agg_accumulator_empty_returns_null() {
996        let acc = StringAggAccumulator::new(",".into());
997        assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
998    }
999
1000    #[test]
1001    fn encode_group_key_is_deterministic() {
1002        let values = vec![
1003            SqlValue::Integer(1),
1004            SqlValue::Text("a".into()),
1005            SqlValue::Null,
1006        ];
1007        let first = encode_group_key(&values).unwrap();
1008        let second = encode_group_key(&values).unwrap();
1009        assert_eq!(first, second);
1010    }
1011}
alopex_sql/executor/query/aggregate.rs

alopex_sql/executor/query/
aggregate.rs
