Skip to main content

alopex_sql/executor/query/
aggregate.rs

1use std::cmp::Ordering;
2use std::collections::{HashMap, HashSet};
3use std::sync::Arc;
4use std::sync::atomic::{AtomicU64, AtomicUsize};
5
6use crate::catalog::ColumnMetadata;
7use crate::executor::evaluator::EvalContext;
8use crate::executor::memory::{MemoryPolicy, MemoryTracker, map_core_memory_error};
9use crate::executor::{EvaluationError, ExecutorError, Result};
10use crate::planner::aggregate_expr::{AggregateExpr, AggregateFunction};
11use crate::planner::typed_expr::TypedExpr;
12use crate::planner::types::ResolvedType;
13use crate::storage::SqlValue;
14use alopex_core::sql::stream::ByteSized;
15
16use super::{Row, RowIterator, iterator::VecIterator};
17
18/// Byte-encoded group key for hash-based aggregation.
19pub type GroupKeyBytes = Vec<u8>;
20
21fn encode_group_value(value: &SqlValue, buf: &mut Vec<u8>) -> Result<()> {
22    buf.push(value.type_tag());
23    match value {
24        SqlValue::Null => Ok(()),
25        SqlValue::Integer(v) => {
26            buf.extend_from_slice(&v.to_le_bytes());
27            Ok(())
28        }
29        SqlValue::BigInt(v) => {
30            buf.extend_from_slice(&v.to_le_bytes());
31            Ok(())
32        }
33        SqlValue::Float(v) => {
34            buf.extend_from_slice(&v.to_bits().to_le_bytes());
35            Ok(())
36        }
37        SqlValue::Double(v) => {
38            buf.extend_from_slice(&v.to_bits().to_le_bytes());
39            Ok(())
40        }
41        SqlValue::Text(s) => {
42            let len = u32::try_from(s.len()).map_err(|_| ExecutorError::InvalidOperation {
43                operation: "aggregate".into(),
44                reason: "text length exceeds u32::MAX".into(),
45            })?;
46            buf.extend_from_slice(&len.to_le_bytes());
47            buf.extend_from_slice(s.as_bytes());
48            Ok(())
49        }
50        SqlValue::Blob(bytes) => {
51            let len = u32::try_from(bytes.len()).map_err(|_| ExecutorError::InvalidOperation {
52                operation: "aggregate".into(),
53                reason: "blob length exceeds u32::MAX".into(),
54            })?;
55            buf.extend_from_slice(&len.to_le_bytes());
56            buf.extend_from_slice(bytes);
57            Ok(())
58        }
59        SqlValue::Boolean(b) => {
60            buf.push(u8::from(*b));
61            Ok(())
62        }
63        SqlValue::Timestamp(v) => {
64            buf.extend_from_slice(&v.to_le_bytes());
65            Ok(())
66        }
67        SqlValue::Vector(values) => {
68            let len = u32::try_from(values.len()).map_err(|_| ExecutorError::InvalidOperation {
69                operation: "aggregate".into(),
70                reason: "vector length exceeds u32::MAX".into(),
71            })?;
72            buf.extend_from_slice(&len.to_le_bytes());
73            for f in values {
74                buf.extend_from_slice(&f.to_bits().to_le_bytes());
75            }
76            Ok(())
77        }
78    }
79}
80
81/// Encode group key values into a deterministic byte sequence.
82pub fn encode_group_key(values: &[SqlValue]) -> Result<GroupKeyBytes> {
83    let mut buf = Vec::new();
84    for value in values {
85        encode_group_value(value, &mut buf)?;
86    }
87    Ok(buf)
88}
89
90/// Accumulator interface for aggregate function execution.
91pub trait Accumulator: Send {
92    /// Update the accumulator with a new value (None for COUNT(*) rows).
93    fn update(&mut self, value: Option<SqlValue>) -> Result<()>;
94    /// Return the serializable partial aggregate state.
95    fn state(&self) -> Result<Vec<SqlValue>>;
96    /// Merge a partial state produced by an accumulator of the same function.
97    fn merge(&mut self, state: &[SqlValue]) -> Result<()>;
98    /// Finalize the accumulator and return the resulting SqlValue.
99    fn finalize(&self) -> Result<SqlValue>;
100    /// Clone the accumulator as a trait object.
101    fn clone_box(&self) -> Box<dyn Accumulator>;
102}
103
104impl Clone for Box<dyn Accumulator> {
105    fn clone(&self) -> Self {
106        self.clone_box()
107    }
108}
109
110fn invalid_aggregate_state(function: &str, reason: impl Into<String>) -> ExecutorError {
111    ExecutorError::InvalidOperation {
112        operation: function.into(),
113        reason: reason.into(),
114    }
115}
116
117fn expect_state_arity(function: &str, state: &[SqlValue], expected: usize) -> Result<()> {
118    if state.len() == expected {
119        Ok(())
120    } else {
121        Err(invalid_aggregate_state(
122            function,
123            format!("expected {expected} state value(s), got {}", state.len()),
124        ))
125    }
126}
127
128fn state_bigint(function: &str, value: &SqlValue, index: usize) -> Result<i64> {
129    match value {
130        SqlValue::BigInt(v) => Ok(*v),
131        other => Err(invalid_aggregate_state(
132            function,
133            format!(
134                "state value {index} expected BigInt, got {}",
135                other.type_name()
136            ),
137        )),
138    }
139}
140
141fn state_double(function: &str, value: &SqlValue, index: usize) -> Result<f64> {
142    match value {
143        SqlValue::Double(v) => Ok(*v),
144        other => Err(invalid_aggregate_state(
145            function,
146            format!(
147                "state value {index} expected Double, got {}",
148                other.type_name()
149            ),
150        )),
151    }
152}
153
154fn state_text<'a>(function: &str, value: &'a SqlValue, index: usize) -> Result<&'a str> {
155    match value {
156        SqlValue::Text(v) => Ok(v),
157        other => Err(invalid_aggregate_state(
158            function,
159            format!(
160                "state value {index} expected Text, got {}",
161                other.type_name()
162            ),
163        )),
164    }
165}
166
167fn distinct_allows(
168    distinct_values: &mut Option<HashSet<Vec<u8>>>,
169    value: &SqlValue,
170) -> Result<bool> {
171    if value.is_null() {
172        return Ok(false);
173    }
174    let Some(distinct) = distinct_values else {
175        return Ok(true);
176    };
177    let encoded = encode_group_key(std::slice::from_ref(value))?;
178    Ok(distinct.insert(encoded))
179}
180
181/// Accumulator for COUNT / COUNT(DISTINCT).
182#[derive(Debug, Clone)]
183pub struct CountAccumulator {
184    count: usize,
185    distinct_values: Option<HashSet<Vec<u8>>>,
186}
187
188impl CountAccumulator {
189    /// Create a new count accumulator.
190    pub fn new(distinct: bool) -> Self {
191        Self {
192            count: 0,
193            distinct_values: if distinct { Some(HashSet::new()) } else { None },
194        }
195    }
196}
197
198impl Accumulator for CountAccumulator {
199    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
200        match (&mut self.distinct_values, value) {
201            (Some(distinct), Some(value)) => {
202                if value.is_null() {
203                    return Ok(());
204                }
205                let encoded = encode_group_key(std::slice::from_ref(&value))?;
206                if distinct.insert(encoded) {
207                    self.count += 1;
208                }
209            }
210            (Some(_), None) => {
211                self.count += 1;
212            }
213            (None, Some(value)) => {
214                if !value.is_null() {
215                    self.count += 1;
216                }
217            }
218            (None, None) => {
219                self.count += 1;
220            }
221        }
222        Ok(())
223    }
224
225    fn finalize(&self) -> Result<SqlValue> {
226        Ok(SqlValue::BigInt(self.count as i64))
227    }
228
229    fn state(&self) -> Result<Vec<SqlValue>> {
230        Ok(vec![SqlValue::BigInt(self.count as i64)])
231    }
232
233    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
234        expect_state_arity("count", state, 1)?;
235        let count = state_bigint("count", &state[0], 0)?;
236        if count < 0 {
237            return Err(invalid_aggregate_state(
238                "count",
239                "state count must be non-negative",
240            ));
241        }
242        self.count = self.count.saturating_add(count as usize);
243        Ok(())
244    }
245
246    fn clone_box(&self) -> Box<dyn Accumulator> {
247        Box::new(self.clone())
248    }
249}
250
251/// Accumulator for SUM.
252#[derive(Debug, Clone)]
253pub struct SumAccumulator {
254    sum: Option<SqlValue>,
255    result_type: ResolvedType,
256    distinct_values: Option<HashSet<Vec<u8>>>,
257}
258
259impl SumAccumulator {
260    /// Create a new sum accumulator.
261    pub fn new() -> Self {
262        Self::with_distinct(false)
263    }
264
265    pub fn with_distinct(distinct: bool) -> Self {
266        Self::with_distinct_for_type(distinct, ResolvedType::Double)
267    }
268
269    pub fn with_distinct_for_type(distinct: bool, result_type: ResolvedType) -> Self {
270        Self {
271            sum: None,
272            result_type,
273            distinct_values: if distinct { Some(HashSet::new()) } else { None },
274        }
275    }
276
277    fn add_value(&mut self, value: SqlValue) -> Result<()> {
278        let next = match &self.result_type {
279            ResolvedType::Integer => {
280                let SqlValue::Integer(value) = value else {
281                    return sum_type_mismatch("Integer", &value);
282                };
283                let sum = match self.sum.as_ref() {
284                    None => value,
285                    Some(SqlValue::Integer(current)) => current
286                        .checked_add(value)
287                        .ok_or(ExecutorError::Evaluation(EvaluationError::Overflow))?,
288                    Some(other) => return sum_type_mismatch("Integer", other),
289                };
290                SqlValue::Integer(sum)
291            }
292            ResolvedType::BigInt => {
293                let value = match value {
294                    SqlValue::Integer(value) => i64::from(value),
295                    SqlValue::BigInt(value) => value,
296                    other => return sum_type_mismatch("BigInt", &other),
297                };
298                let sum = match self.sum.as_ref() {
299                    None => value,
300                    Some(SqlValue::BigInt(current)) => current
301                        .checked_add(value)
302                        .ok_or(ExecutorError::Evaluation(EvaluationError::Overflow))?,
303                    Some(other) => return sum_type_mismatch("BigInt", other),
304                };
305                SqlValue::BigInt(sum)
306            }
307            _ => {
308                let value = numeric_to_f64(&value)?;
309                let sum = match self.sum.as_ref() {
310                    None => value,
311                    Some(SqlValue::Double(current)) => *current + value,
312                    Some(other) => return sum_type_mismatch("Double", other),
313                };
314                SqlValue::Double(sum)
315            }
316        };
317        self.sum = Some(next);
318        Ok(())
319    }
320}
321
322impl Default for SumAccumulator {
323    fn default() -> Self {
324        Self::new()
325    }
326}
327
328impl Accumulator for SumAccumulator {
329    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
330        let Some(value) = value else {
331            return Ok(());
332        };
333        if value.is_null() {
334            return Ok(());
335        }
336        if !distinct_allows(&mut self.distinct_values, &value)? {
337            return Ok(());
338        }
339        self.add_value(value)
340    }
341
342    fn finalize(&self) -> Result<SqlValue> {
343        Ok(self.sum.clone().unwrap_or(SqlValue::Null))
344    }
345
346    fn state(&self) -> Result<Vec<SqlValue>> {
347        Ok(vec![self.sum.clone().unwrap_or(SqlValue::Null)])
348    }
349
350    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
351        expect_state_arity("sum", state, 1)?;
352        if state[0].is_null() {
353            return Ok(());
354        }
355        self.add_value(state[0].clone())
356    }
357
358    fn clone_box(&self) -> Box<dyn Accumulator> {
359        Box::new(self.clone())
360    }
361}
362
363/// Accumulator for TOTAL (SUM that returns 0.0 on empty/all-NULL input).
364#[derive(Debug, Clone)]
365pub struct TotalAccumulator {
366    sum: Option<f64>,
367}
368
369impl TotalAccumulator {
370    /// Create a new total accumulator.
371    pub fn new() -> Self {
372        Self { sum: None }
373    }
374}
375
376impl Default for TotalAccumulator {
377    fn default() -> Self {
378        Self::new()
379    }
380}
381
382impl Accumulator for TotalAccumulator {
383    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
384        let Some(value) = value else {
385            return Ok(());
386        };
387        if value.is_null() {
388            return Ok(());
389        }
390        let numeric = numeric_to_f64(&value)?;
391        self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
392        Ok(())
393    }
394
395    fn finalize(&self) -> Result<SqlValue> {
396        Ok(SqlValue::Double(self.sum.unwrap_or(0.0)))
397    }
398
399    fn state(&self) -> Result<Vec<SqlValue>> {
400        Ok(vec![SqlValue::Double(self.sum.unwrap_or(0.0))])
401    }
402
403    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
404        expect_state_arity("total", state, 1)?;
405        let value = state_double("total", &state[0], 0)?;
406        self.sum = Some(self.sum.unwrap_or(0.0) + value);
407        Ok(())
408    }
409
410    fn clone_box(&self) -> Box<dyn Accumulator> {
411        Box::new(self.clone())
412    }
413}
414
415/// Accumulator for AVG.
416#[derive(Debug, Clone)]
417pub struct AvgAccumulator {
418    sum: Option<f64>,
419    count: usize,
420    distinct_values: Option<HashSet<Vec<u8>>>,
421}
422
423impl AvgAccumulator {
424    /// Create a new average accumulator.
425    pub fn new() -> Self {
426        Self::with_distinct(false)
427    }
428
429    pub fn with_distinct(distinct: bool) -> Self {
430        Self {
431            sum: None,
432            count: 0,
433            distinct_values: if distinct { Some(HashSet::new()) } else { None },
434        }
435    }
436}
437
438impl Default for AvgAccumulator {
439    fn default() -> Self {
440        Self::new()
441    }
442}
443
444impl Accumulator for AvgAccumulator {
445    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
446        let Some(value) = value else {
447            return Ok(());
448        };
449        if value.is_null() {
450            return Ok(());
451        }
452        if !distinct_allows(&mut self.distinct_values, &value)? {
453            return Ok(());
454        }
455        let numeric = numeric_to_f64(&value)?;
456        self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
457        self.count += 1;
458        Ok(())
459    }
460
461    fn finalize(&self) -> Result<SqlValue> {
462        if self.count == 0 {
463            return Ok(SqlValue::Null);
464        }
465        let sum = self.sum.unwrap_or(0.0);
466        Ok(SqlValue::Double(sum / self.count as f64))
467    }
468
469    fn state(&self) -> Result<Vec<SqlValue>> {
470        Ok(vec![
471            SqlValue::Double(self.sum.unwrap_or(0.0)),
472            SqlValue::BigInt(self.count as i64),
473        ])
474    }
475
476    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
477        expect_state_arity("avg", state, 2)?;
478        let sum = state_double("avg", &state[0], 0)?;
479        let count = state_bigint("avg", &state[1], 1)?;
480        if count < 0 {
481            return Err(invalid_aggregate_state(
482                "avg",
483                "state count must be non-negative",
484            ));
485        }
486        self.sum = Some(self.sum.unwrap_or(0.0) + sum);
487        self.count = self.count.saturating_add(count as usize);
488        Ok(())
489    }
490
491    fn clone_box(&self) -> Box<dyn Accumulator> {
492        Box::new(self.clone())
493    }
494}
495
496fn numeric_to_f64(value: &SqlValue) -> Result<f64> {
497    match value {
498        SqlValue::Integer(v) => Ok(*v as f64),
499        SqlValue::BigInt(v) => Ok(*v as f64),
500        SqlValue::Float(v) => Ok(*v as f64),
501        SqlValue::Double(v) => Ok(*v),
502        _ => Err(ExecutorError::Evaluation(
503            crate::executor::EvaluationError::TypeMismatch {
504                expected: "numeric".into(),
505                actual: value.type_name().into(),
506            },
507        )),
508    }
509}
510
511fn sum_type_mismatch<T>(expected: &str, actual: &SqlValue) -> Result<T> {
512    Err(ExecutorError::Evaluation(EvaluationError::TypeMismatch {
513        expected: expected.into(),
514        actual: actual.type_name().into(),
515    }))
516}
517
518/// Accumulator for MIN / MAX.
519#[derive(Debug, Clone)]
520pub struct MinMaxAccumulator {
521    value: Option<SqlValue>,
522    is_min: bool,
523    distinct_values: Option<HashSet<Vec<u8>>>,
524}
525
526impl MinMaxAccumulator {
527    /// Create a new min/max accumulator.
528    pub fn new(is_min: bool) -> Self {
529        Self::with_distinct(is_min, false)
530    }
531
532    pub fn with_distinct(is_min: bool, distinct: bool) -> Self {
533        Self {
534            value: None,
535            is_min,
536            distinct_values: if distinct { Some(HashSet::new()) } else { None },
537        }
538    }
539}
540
541impl Accumulator for MinMaxAccumulator {
542    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
543        let Some(value) = value else {
544            return Ok(());
545        };
546        if value.is_null() {
547            return Ok(());
548        }
549        if !distinct_allows(&mut self.distinct_values, &value)? {
550            return Ok(());
551        }
552
553        match &self.value {
554            None => {
555                self.value = Some(value);
556            }
557            Some(current) => {
558                if std::mem::discriminant(current) != std::mem::discriminant(&value) {
559                    return Err(ExecutorError::Evaluation(
560                        crate::executor::EvaluationError::TypeMismatch {
561                            expected: current.type_name().into(),
562                            actual: value.type_name().into(),
563                        },
564                    ));
565                }
566                let ordering = value.partial_cmp(current).ok_or_else(|| {
567                    ExecutorError::Evaluation(crate::executor::EvaluationError::TypeMismatch {
568                        expected: current.type_name().into(),
569                        actual: value.type_name().into(),
570                    })
571                })?;
572                let should_replace = matches!(
573                    (self.is_min, ordering),
574                    (true, Ordering::Less) | (false, Ordering::Greater)
575                );
576                if should_replace {
577                    self.value = Some(value);
578                }
579            }
580        }
581        Ok(())
582    }
583
584    fn finalize(&self) -> Result<SqlValue> {
585        Ok(self.value.clone().unwrap_or(SqlValue::Null))
586    }
587
588    fn state(&self) -> Result<Vec<SqlValue>> {
589        Ok(vec![self.value.clone().unwrap_or(SqlValue::Null)])
590    }
591
592    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
593        expect_state_arity(if self.is_min { "min" } else { "max" }, state, 1)?;
594        if state[0].is_null() {
595            return Ok(());
596        }
597        self.update(Some(state[0].clone()))
598    }
599
600    fn clone_box(&self) -> Box<dyn Accumulator> {
601        Box::new(self.clone())
602    }
603}
604
605/// Accumulator for GROUP_CONCAT.
606#[derive(Debug, Clone)]
607pub struct GroupConcatAccumulator {
608    values: Vec<String>,
609    separator: String,
610    distinct_values: Option<HashSet<Vec<u8>>>,
611}
612
613impl GroupConcatAccumulator {
614    /// Create a new GROUP_CONCAT accumulator with the given separator.
615    pub fn new(separator: String) -> Self {
616        Self::with_distinct(separator, false)
617    }
618
619    pub fn with_distinct(separator: String, distinct: bool) -> Self {
620        Self {
621            values: Vec::new(),
622            separator,
623            distinct_values: if distinct { Some(HashSet::new()) } else { None },
624        }
625    }
626}
627
628impl Accumulator for GroupConcatAccumulator {
629    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
630        let Some(value) = value else {
631            return Ok(());
632        };
633        match value {
634            SqlValue::Null => Ok(()),
635            SqlValue::Text(text) => {
636                let value = SqlValue::Text(text.clone());
637                if !distinct_allows(&mut self.distinct_values, &value)? {
638                    return Ok(());
639                }
640                self.values.push(text);
641                Ok(())
642            }
643            other => Err(ExecutorError::Evaluation(
644                crate::executor::EvaluationError::TypeMismatch {
645                    expected: "Text".into(),
646                    actual: other.type_name().into(),
647                },
648            )),
649        }
650    }
651
652    fn finalize(&self) -> Result<SqlValue> {
653        if self.values.is_empty() {
654            return Ok(SqlValue::Null);
655        }
656        Ok(SqlValue::Text(self.values.join(&self.separator)))
657    }
658
659    fn state(&self) -> Result<Vec<SqlValue>> {
660        Ok(vec![
661            if self.values.is_empty() {
662                SqlValue::Null
663            } else {
664                SqlValue::Text(self.values.join(&self.separator))
665            },
666            SqlValue::Text(self.separator.clone()),
667        ])
668    }
669
670    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
671        expect_state_arity("group_concat", state, 2)?;
672        let separator = state_text("group_concat", &state[1], 1)?;
673        if separator != self.separator {
674            return Err(invalid_aggregate_state(
675                "group_concat",
676                "state separator differs from accumulator separator",
677            ));
678        }
679        match &state[0] {
680            SqlValue::Null => Ok(()),
681            SqlValue::Text(text) => {
682                self.values.push(text.clone());
683                Ok(())
684            }
685            other => Err(invalid_aggregate_state(
686                "group_concat",
687                format!(
688                    "state value 0 expected Text or Null, got {}",
689                    other.type_name()
690                ),
691            )),
692        }
693    }
694
695    fn clone_box(&self) -> Box<dyn Accumulator> {
696        Box::new(self.clone())
697    }
698}
699
700/// Accumulator for STRING_AGG.
701#[derive(Debug, Clone)]
702pub struct StringAggAccumulator {
703    values: Vec<String>,
704    separator: String,
705    distinct_values: Option<HashSet<Vec<u8>>>,
706}
707
708impl StringAggAccumulator {
709    /// Create a new string_agg accumulator.
710    pub fn new(separator: String) -> Self {
711        Self::with_distinct(separator, false)
712    }
713
714    pub fn with_distinct(separator: String, distinct: bool) -> Self {
715        Self {
716            values: Vec::new(),
717            separator,
718            distinct_values: if distinct { Some(HashSet::new()) } else { None },
719        }
720    }
721}
722
723impl Accumulator for StringAggAccumulator {
724    fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
725        let Some(value) = value else {
726            return Ok(());
727        };
728        match value {
729            SqlValue::Null => Ok(()),
730            SqlValue::Text(s) => {
731                let value = SqlValue::Text(s.clone());
732                if !distinct_allows(&mut self.distinct_values, &value)? {
733                    return Ok(());
734                }
735                self.values.push(s);
736                Ok(())
737            }
738            other => Err(ExecutorError::Evaluation(
739                crate::executor::EvaluationError::TypeMismatch {
740                    expected: "Text".into(),
741                    actual: other.type_name().into(),
742                },
743            )),
744        }
745    }
746
747    fn finalize(&self) -> Result<SqlValue> {
748        if self.values.is_empty() {
749            return Ok(SqlValue::Null);
750        }
751        Ok(SqlValue::Text(self.values.join(&self.separator)))
752    }
753
754    fn state(&self) -> Result<Vec<SqlValue>> {
755        Ok(vec![
756            if self.values.is_empty() {
757                SqlValue::Null
758            } else {
759                SqlValue::Text(self.values.join(&self.separator))
760            },
761            SqlValue::Text(self.separator.clone()),
762        ])
763    }
764
765    fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
766        expect_state_arity("string_agg", state, 2)?;
767        let separator = state_text("string_agg", &state[1], 1)?;
768        if separator != self.separator {
769            return Err(invalid_aggregate_state(
770                "string_agg",
771                "state separator differs from accumulator separator",
772            ));
773        }
774        match &state[0] {
775            SqlValue::Null => Ok(()),
776            SqlValue::Text(text) => {
777                self.values.push(text.clone());
778                Ok(())
779            }
780            other => Err(invalid_aggregate_state(
781                "string_agg",
782                format!(
783                    "state value 0 expected Text or Null, got {}",
784                    other.type_name()
785                ),
786            )),
787        }
788    }
789
790    fn clone_box(&self) -> Box<dyn Accumulator> {
791        Box::new(self.clone())
792    }
793}
794
795/// Create a new accumulator instance for the aggregate function.
796pub fn create_accumulator(function: &AggregateFunction, distinct: bool) -> Box<dyn Accumulator> {
797    match function {
798        AggregateFunction::Count => Box::new(CountAccumulator::new(distinct)),
799        AggregateFunction::Sum => Box::new(SumAccumulator::with_distinct(distinct)),
800        AggregateFunction::Total => Box::new(TotalAccumulator::new()),
801        AggregateFunction::Avg => Box::new(AvgAccumulator::with_distinct(distinct)),
802        AggregateFunction::Min => Box::new(MinMaxAccumulator::with_distinct(true, distinct)),
803        AggregateFunction::Max => Box::new(MinMaxAccumulator::with_distinct(false, distinct)),
804        AggregateFunction::GroupConcat { separator } => {
805            let sep = separator.clone().unwrap_or_else(|| ",".to_string());
806            Box::new(GroupConcatAccumulator::with_distinct(sep, distinct))
807        }
808        AggregateFunction::StringAgg { separator } => {
809            let sep = separator.clone().unwrap_or_else(|| ",".to_string());
810            Box::new(StringAggAccumulator::with_distinct(sep, distinct))
811        }
812    }
813}
814
815/// Create an accumulator using the aggregate expression's resolved result type.
816pub fn create_accumulator_for_aggregate(aggregate: &AggregateExpr) -> Box<dyn Accumulator> {
817    match &aggregate.function {
818        AggregateFunction::Sum => Box::new(SumAccumulator::with_distinct_for_type(
819            aggregate.distinct,
820            aggregate.result_type.clone(),
821        )),
822        _ => create_accumulator(&aggregate.function, aggregate.distinct),
823    }
824}
825
826/// Returns whether an aggregate's partial state can be merged without
827/// changing the current local SQL result. Floating-point, DISTINCT, and
828/// order-sensitive aggregates must instead be replayed from ordered inputs.
829pub fn exact_partial_aggregate_is_proven(aggregate: &AggregateExpr) -> bool {
830    !aggregate.distinct
831        && matches!(
832            aggregate.function,
833            AggregateFunction::Count | AggregateFunction::Min | AggregateFunction::Max
834        )
835}
836
837/// Merge only aggregate states whose merge rule is proven to preserve the
838/// current local SQL result. This is the coordinator-side kernel used by the
839/// distributed result assembler after every worker has acknowledged cleanup.
840pub fn merge_exact_aggregate_states(
841    aggregates: &[AggregateExpr],
842    partial_rows: impl IntoIterator<Item = Vec<Vec<SqlValue>>>,
843) -> Result<Vec<SqlValue>> {
844    if let Some(aggregate) = aggregates
845        .iter()
846        .find(|aggregate| !exact_partial_aggregate_is_proven(aggregate))
847    {
848        return Err(ExecutorError::InvalidOperation {
849            operation: "distributed aggregate merge".into(),
850            reason: format!(
851                "{:?} requires ordered input replay rather than an exact partial merge",
852                aggregate.function
853            ),
854        });
855    }
856
857    let mut accumulators = aggregates
858        .iter()
859        .map(create_accumulator_for_aggregate)
860        .collect::<Vec<_>>();
861    for states in partial_rows {
862        if states.len() != accumulators.len() {
863            return Err(ExecutorError::InvalidOperation {
864                operation: "distributed aggregate merge".into(),
865                reason: format!(
866                    "partial state has {} aggregate(s), expected {}",
867                    states.len(),
868                    accumulators.len()
869                ),
870            });
871        }
872        for (accumulator, state) in accumulators.iter_mut().zip(states) {
873            accumulator.merge(&state)?;
874        }
875    }
876    accumulators
877        .iter()
878        .map(|accumulator| accumulator.finalize())
879        .collect()
880}
881
882const DEFAULT_GROUP_LIMIT: usize = 1_000_000;
883const AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES: u64 = 32;
884
885/// Aggregate execution mode.
886#[derive(Debug, Clone, Copy, PartialEq, Eq)]
887pub enum AggregateMode {
888    /// Consume raw input rows and output aggregate partial state rows.
889    Partial,
890    /// Consume partial state rows and output final aggregate values.
891    Final,
892    /// Consume raw input rows and output final aggregate values in one pass.
893    Single,
894}
895
896struct AggregateGroup {
897    key_values: Vec<SqlValue>,
898    accumulators: Vec<Box<dyn Accumulator>>,
899}
900
901/// Iterator that performs hash-based aggregation over input rows.
902pub struct AggregateIterator<'a> {
903    input: Box<dyn RowIterator + 'a>,
904    group_keys: Vec<TypedExpr>,
905    aggregates: Vec<AggregateExpr>,
906    having: Option<TypedExpr>,
907    mode: AggregateMode,
908    hash_table: Option<HashMap<GroupKeyBytes, AggregateGroup>>,
909    result_rows: Vec<Row>,
910    index: usize,
911    schema: Vec<ColumnMetadata>,
912    group_limit: usize,
913    memory_tracker: Option<MemoryTracker>,
914    shared_group_counter: Option<Arc<AtomicUsize>>,
915}
916
917impl<'a> AggregateIterator<'a> {
918    /// Create a new aggregate iterator with the default group limit.
919    pub fn new(
920        input: Box<dyn RowIterator + 'a>,
921        group_keys: Vec<TypedExpr>,
922        aggregates: Vec<AggregateExpr>,
923        having: Option<TypedExpr>,
924        schema: Vec<ColumnMetadata>,
925    ) -> Self {
926        Self {
927            input,
928            group_keys,
929            aggregates,
930            having,
931            mode: AggregateMode::Single,
932            hash_table: None,
933            result_rows: Vec::new(),
934            index: 0,
935            schema,
936            group_limit: DEFAULT_GROUP_LIMIT,
937            memory_tracker: None,
938            shared_group_counter: None,
939        }
940    }
941
942    /// Override the maximum number of groups allowed during aggregation.
943    pub fn with_group_limit(mut self, limit: usize) -> Self {
944        self.group_limit = limit;
945        self
946    }
947
948    /// Set the aggregate execution mode.
949    pub fn with_mode(mut self, mode: AggregateMode) -> Self {
950        self.mode = mode;
951        self
952    }
953
954    /// Attach a memory policy for enforcing in-flight aggregation limits.
955    pub fn with_memory_policy(mut self, policy: Option<MemoryPolicy>) -> Self {
956        self.memory_tracker = policy.map(MemoryTracker::new);
957        self
958    }
959
960    /// Attach a shared group counter used by parallel Partial aggregation.
961    pub fn with_shared_group_counter(mut self, counter: Option<Arc<AtomicUsize>>) -> Self {
962        self.shared_group_counter = counter;
963        self
964    }
965
966    fn build_hash_table(&mut self) -> Result<()> {
967        let mut table: HashMap<GroupKeyBytes, AggregateGroup> = HashMap::new();
968        let mut next_row_id = 0u64;
969
970        while let Some(result) = self.input.next_row() {
971            let row = result?;
972            let (key_values, key_bytes) = match self.mode {
973                AggregateMode::Final => {
974                    let key_values = row
975                        .values
976                        .get(..self.group_keys.len())
977                        .ok_or_else(|| {
978                            invalid_aggregate_state(
979                                "aggregate",
980                                "partial state row is missing group key values",
981                            )
982                        })?
983                        .to_vec();
984                    let key_bytes = encode_group_key(&key_values)?;
985                    (key_values, key_bytes)
986                }
987                AggregateMode::Partial | AggregateMode::Single => {
988                    let ctx = EvalContext::new(&row.values);
989                    let mut key_values = Vec::with_capacity(self.group_keys.len());
990                    for expr in &self.group_keys {
991                        key_values.push(crate::executor::evaluator::evaluate(expr, &ctx)?);
992                    }
993                    let key_bytes = encode_group_key(&key_values)?;
994                    (key_values, key_bytes)
995                }
996            };
997
998            if !table.contains_key(&key_bytes) {
999                self.reserve_group_slot(table.len())?;
1000                if let Some(tracker) = &mut self.memory_tracker {
1001                    tracker
1002                        .add_values(&key_values)
1003                        .map_err(map_core_memory_error)?;
1004                    tracker
1005                        .add_bytes(
1006                            self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES,
1007                        )
1008                        .map_err(map_core_memory_error)?;
1009                }
1010                let accumulators = self
1011                    .aggregates
1012                    .iter()
1013                    .map(|agg| {
1014                        let mut aggregate = agg.clone();
1015                        aggregate.distinct =
1016                            matches!(self.mode, AggregateMode::Single) && agg.distinct;
1017                        create_accumulator_for_aggregate(&aggregate)
1018                    })
1019                    .collect::<Vec<_>>();
1020                table.insert(
1021                    key_bytes.clone(),
1022                    AggregateGroup {
1023                        key_values: key_values.clone(),
1024                        accumulators,
1025                    },
1026                );
1027            }
1028
1029            if let Some(group) = table.get_mut(&key_bytes) {
1030                match self.mode {
1031                    AggregateMode::Final => {
1032                        let mut offset = self.group_keys.len();
1033                        for (idx, agg) in self.aggregates.iter().enumerate() {
1034                            let arity = aggregate_state_types(agg).len();
1035                            let state = row.values.get(offset..offset + arity).ok_or_else(|| {
1036                                invalid_aggregate_state(
1037                                    "aggregate",
1038                                    format!(
1039                                        "partial state row is missing state values for aggregate {idx}"
1040                                    ),
1041                                )
1042                            })?;
1043                            group.accumulators[idx].merge(state)?;
1044                            offset += arity;
1045                        }
1046                        if offset != row.values.len() {
1047                            return Err(invalid_aggregate_state(
1048                                "aggregate",
1049                                format!(
1050                                    "partial state row has {} trailing value(s)",
1051                                    row.values.len() - offset
1052                                ),
1053                            ));
1054                        }
1055                    }
1056                    AggregateMode::Partial | AggregateMode::Single => {
1057                        let ctx = EvalContext::new(&row.values);
1058                        for (idx, agg) in self.aggregates.iter().enumerate() {
1059                            let value = match &agg.arg {
1060                                None => None,
1061                                Some(expr) => {
1062                                    Some(crate::executor::evaluator::evaluate(expr, &ctx)?)
1063                                }
1064                            };
1065                            if let Some(tracker) = &mut self.memory_tracker
1066                                && matches!(
1067                                    agg.function,
1068                                    AggregateFunction::GroupConcat { .. }
1069                                        | AggregateFunction::StringAgg { .. }
1070                                )
1071                                && let Some(value_ref) = value.as_ref()
1072                            {
1073                                tracker
1074                                    .add_value(value_ref)
1075                                    .map_err(map_core_memory_error)?;
1076                            }
1077                            group.accumulators[idx].update(value)?;
1078                        }
1079                    }
1080                }
1081            }
1082        }
1083
1084        if table.is_empty() && self.group_keys.is_empty() {
1085            if let Some(tracker) = &mut self.memory_tracker {
1086                tracker
1087                    .add_bytes(self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES)
1088                    .map_err(map_core_memory_error)?;
1089            }
1090            let accumulators = self
1091                .aggregates
1092                .iter()
1093                .map(|agg| {
1094                    let mut aggregate = agg.clone();
1095                    aggregate.distinct = matches!(self.mode, AggregateMode::Single) && agg.distinct;
1096                    create_accumulator_for_aggregate(&aggregate)
1097                })
1098                .collect::<Vec<_>>();
1099            table.insert(
1100                Vec::new(),
1101                AggregateGroup {
1102                    key_values: Vec::new(),
1103                    accumulators,
1104                },
1105            );
1106        }
1107
1108        let mut rows = Vec::with_capacity(table.len());
1109        for group in table.values() {
1110            let mut values = Vec::with_capacity(self.group_keys.len() + self.aggregates.len());
1111            values.extend(group.key_values.iter().cloned());
1112            for acc in &group.accumulators {
1113                match self.mode {
1114                    AggregateMode::Partial => values.extend(acc.state()?),
1115                    AggregateMode::Final | AggregateMode::Single => values.push(acc.finalize()?),
1116                }
1117            }
1118            let row = Row::new(next_row_id, values);
1119            next_row_id += 1;
1120            if let Some(tracker) = &mut self.memory_tracker {
1121                tracker
1122                    .add_row(&row.values)
1123                    .map_err(map_core_memory_error)?;
1124            }
1125
1126            if self.mode != AggregateMode::Partial
1127                && let Some(having) = &self.having
1128            {
1129                let ctx = EvalContext::new(&row.values);
1130                match crate::executor::evaluator::evaluate(having, &ctx)? {
1131                    SqlValue::Boolean(true) => rows.push(row),
1132                    SqlValue::Boolean(false) | SqlValue::Null => {}
1133                    other => {
1134                        return Err(ExecutorError::Evaluation(
1135                            crate::executor::EvaluationError::TypeMismatch {
1136                                expected: "Boolean".into(),
1137                                actual: other.type_name().into(),
1138                            },
1139                        ));
1140                    }
1141                }
1142            } else {
1143                rows.push(row);
1144            }
1145        }
1146
1147        self.hash_table = Some(table);
1148        self.result_rows = rows;
1149        Ok(())
1150    }
1151
1152    fn reserve_group_slot(&self, local_group_count: usize) -> Result<()> {
1153        let next_count = if let Some(counter) = &self.shared_group_counter {
1154            counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1
1155        } else {
1156            local_group_count + 1
1157        };
1158        if next_count > self.group_limit {
1159            return Err(ExecutorError::ResourceExhausted {
1160                message: format!(
1161                    "GROUP BY result exceeds memory limit (max groups: {})",
1162                    self.group_limit
1163                ),
1164            });
1165        }
1166        Ok(())
1167    }
1168}
1169
1170impl<'a> RowIterator for AggregateIterator<'a> {
1171    fn next_row(&mut self) -> Option<Result<Row>> {
1172        if self.hash_table.is_none()
1173            && let Err(err) = self.build_hash_table()
1174        {
1175            return Some(Err(err));
1176        }
1177
1178        if self.index >= self.result_rows.len() {
1179            return None;
1180        }
1181        let row = self.result_rows[self.index].clone();
1182        self.index += 1;
1183        Some(Ok(row))
1184    }
1185
1186    fn schema(&self) -> &[ColumnMetadata] {
1187        &self.schema
1188    }
1189}
1190
1191/// Iterator that performs streaming aggregation over sorted input.
1192pub struct StreamingAggregateIterator<'a> {
1193    input: Box<dyn RowIterator + 'a>,
1194    group_keys: Vec<TypedExpr>,
1195    aggregates: Vec<AggregateExpr>,
1196    having: Option<TypedExpr>,
1197    schema: Vec<ColumnMetadata>,
1198    current_key: Option<Vec<SqlValue>>,
1199    accumulators: Vec<Box<dyn Accumulator>>,
1200    pending_row: Option<Row>,
1201    finished: bool,
1202    next_row_id: u64,
1203    saw_row: bool,
1204}
1205
1206impl<'a> StreamingAggregateIterator<'a> {
1207    pub fn new(
1208        input: Box<dyn RowIterator + 'a>,
1209        group_keys: Vec<TypedExpr>,
1210        aggregates: Vec<AggregateExpr>,
1211        having: Option<TypedExpr>,
1212        schema: Vec<ColumnMetadata>,
1213    ) -> Self {
1214        Self {
1215            input,
1216            group_keys,
1217            aggregates,
1218            having,
1219            schema,
1220            current_key: None,
1221            accumulators: Vec::new(),
1222            pending_row: None,
1223            finished: false,
1224            next_row_id: 0,
1225            saw_row: false,
1226        }
1227    }
1228
1229    fn init_accumulators(&self) -> Vec<Box<dyn Accumulator>> {
1230        self.aggregates
1231            .iter()
1232            .map(create_accumulator_for_aggregate)
1233            .collect()
1234    }
1235
1236    fn update_accumulators(&mut self, ctx: &EvalContext<'_>) -> Result<()> {
1237        for (idx, agg) in self.aggregates.iter().enumerate() {
1238            let value = match &agg.arg {
1239                None => None,
1240                Some(expr) => Some(crate::executor::evaluator::evaluate(expr, ctx)?),
1241            };
1242            self.accumulators[idx].update(value)?;
1243        }
1244        Ok(())
1245    }
1246
1247    fn finalize_group(&mut self, key_values: &[SqlValue]) -> Result<Option<Row>> {
1248        let mut values = Vec::with_capacity(self.group_keys.len() + self.aggregates.len());
1249        values.extend(key_values.iter().cloned());
1250        for acc in &self.accumulators {
1251            values.push(acc.finalize()?);
1252        }
1253        let row = Row::new(self.next_row_id, values);
1254        self.next_row_id = self.next_row_id.saturating_add(1);
1255
1256        if let Some(having) = &self.having {
1257            let ctx = EvalContext::new(&row.values);
1258            match crate::executor::evaluator::evaluate(having, &ctx)? {
1259                SqlValue::Boolean(true) => Ok(Some(row)),
1260                SqlValue::Boolean(false) | SqlValue::Null => Ok(None),
1261                other => Err(ExecutorError::Evaluation(
1262                    crate::executor::EvaluationError::TypeMismatch {
1263                        expected: "Boolean".into(),
1264                        actual: other.type_name().into(),
1265                    },
1266                )),
1267            }
1268        } else {
1269            Ok(Some(row))
1270        }
1271    }
1272}
1273
1274impl<'a> RowIterator for StreamingAggregateIterator<'a> {
1275    fn next_row(&mut self) -> Option<Result<Row>> {
1276        if let Some(row) = self.pending_row.take() {
1277            return Some(Ok(row));
1278        }
1279        if self.finished {
1280            return None;
1281        }
1282
1283        loop {
1284            match self.input.next_row() {
1285                Some(Ok(row)) => {
1286                    self.saw_row = true;
1287                    let ctx = EvalContext::new(&row.values);
1288                    let mut key_values = Vec::with_capacity(self.group_keys.len());
1289                    for expr in &self.group_keys {
1290                        match crate::executor::evaluator::evaluate(expr, &ctx) {
1291                            Ok(value) => key_values.push(value),
1292                            Err(err) => return Some(Err(err)),
1293                        }
1294                    }
1295
1296                    match &self.current_key {
1297                        None => {
1298                            self.current_key = Some(key_values);
1299                            self.accumulators = self.init_accumulators();
1300                            if let Err(err) = self.update_accumulators(&ctx) {
1301                                return Some(Err(err));
1302                            }
1303                        }
1304                        Some(current_key) if *current_key == key_values => {
1305                            if let Err(err) = self.update_accumulators(&ctx) {
1306                                return Some(Err(err));
1307                            }
1308                        }
1309                        Some(_) => {
1310                            let current_key = self.current_key.clone().unwrap_or_default();
1311                            let output = match self.finalize_group(&current_key) {
1312                                Ok(value) => value,
1313                                Err(err) => return Some(Err(err)),
1314                            };
1315                            self.current_key = Some(key_values);
1316                            self.accumulators = self.init_accumulators();
1317                            if let Err(err) = self.update_accumulators(&ctx) {
1318                                return Some(Err(err));
1319                            }
1320                            if let Some(row) = output {
1321                                return Some(Ok(row));
1322                            }
1323                        }
1324                    }
1325                }
1326                Some(Err(err)) => return Some(Err(err)),
1327                None => {
1328                    self.finished = true;
1329                    if let Some(current_key) = self.current_key.take() {
1330                        return match self.finalize_group(&current_key) {
1331                            Ok(Some(row)) => Some(Ok(row)),
1332                            Ok(None) => None,
1333                            Err(err) => Some(Err(err)),
1334                        };
1335                    }
1336
1337                    if self.group_keys.is_empty() && !self.saw_row {
1338                        self.accumulators = self.init_accumulators();
1339                        return match self.finalize_group(&[]) {
1340                            Ok(Some(row)) => Some(Ok(row)),
1341                            Ok(None) => None,
1342                            Err(err) => Some(Err(err)),
1343                        };
1344                    }
1345
1346                    return None;
1347                }
1348            }
1349        }
1350    }
1351
1352    fn schema(&self) -> &[ColumnMetadata] {
1353        &self.schema
1354    }
1355}
1356
1357fn aggregate_state_types(agg: &AggregateExpr) -> Vec<ResolvedType> {
1358    match &agg.function {
1359        AggregateFunction::Count => vec![ResolvedType::BigInt],
1360        AggregateFunction::Sum => vec![agg.result_type.clone()],
1361        AggregateFunction::Total => vec![ResolvedType::Double],
1362        AggregateFunction::Avg => vec![ResolvedType::Double, ResolvedType::BigInt],
1363        AggregateFunction::Min | AggregateFunction::Max => vec![agg.result_type.clone()],
1364        AggregateFunction::GroupConcat { .. } | AggregateFunction::StringAgg { .. } => {
1365            vec![ResolvedType::Text, ResolvedType::Text]
1366        }
1367    }
1368}
1369
1370/// Build output schema for aggregate results.
1371pub fn build_aggregate_schema(
1372    group_keys: &[TypedExpr],
1373    aggregates: &[AggregateExpr],
1374) -> Vec<ColumnMetadata> {
1375    let mut schema = Vec::new();
1376    for (idx, key) in group_keys.iter().enumerate() {
1377        let name = match &key.kind {
1378            crate::planner::typed_expr::TypedExprKind::ColumnRef { column, .. } => column.clone(),
1379            _ => format!("group_{idx}"),
1380        };
1381        schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
1382    }
1383    for (idx, agg) in aggregates.iter().enumerate() {
1384        let name = match &agg.function {
1385            AggregateFunction::Count => format!("count_{idx}"),
1386            AggregateFunction::Sum => format!("sum_{idx}"),
1387            AggregateFunction::Total => format!("total_{idx}"),
1388            AggregateFunction::Avg => format!("avg_{idx}"),
1389            AggregateFunction::Min => format!("min_{idx}"),
1390            AggregateFunction::Max => format!("max_{idx}"),
1391            AggregateFunction::GroupConcat { .. } => format!("group_concat_{idx}"),
1392            AggregateFunction::StringAgg { .. } => format!("string_agg_{idx}"),
1393        };
1394        schema.push(ColumnMetadata::new(name, agg.result_type.clone()));
1395    }
1396    schema
1397}
1398
1399/// Build internal partial aggregate schema.
1400pub fn build_partial_aggregate_schema(
1401    group_keys: &[TypedExpr],
1402    aggregates: &[AggregateExpr],
1403) -> Vec<ColumnMetadata> {
1404    let mut schema = Vec::new();
1405    for (idx, key) in group_keys.iter().enumerate() {
1406        let name = match &key.kind {
1407            crate::planner::typed_expr::TypedExprKind::ColumnRef { column, .. } => column.clone(),
1408            _ => format!("group_{idx}"),
1409        };
1410        schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
1411    }
1412    for (agg_idx, agg) in aggregates.iter().enumerate() {
1413        for (state_idx, state_type) in aggregate_state_types(agg).into_iter().enumerate() {
1414            schema.push(ColumnMetadata::new(
1415                format!("__agg{agg_idx}_state{state_idx}"),
1416                state_type,
1417            ));
1418        }
1419    }
1420    schema
1421}
1422
1423/// Return true when aggregate execution must remain Single for correctness.
1424pub fn should_use_single_for_parallel(parallelism: usize, aggregates: &[AggregateExpr]) -> bool {
1425    parallelism <= 1 || aggregates.iter().any(|agg| agg.distinct)
1426}
1427
1428fn collect_iterator_rows(iter: &mut dyn RowIterator) -> Result<Vec<Row>> {
1429    let mut rows = Vec::new();
1430    while let Some(result) = iter.next_row() {
1431        rows.push(result?);
1432    }
1433    Ok(rows)
1434}
1435
1436struct ChainRowIterator<'a> {
1437    prefix: std::vec::IntoIter<Row>,
1438    tail: Box<dyn RowIterator + 'a>,
1439    schema: Vec<ColumnMetadata>,
1440}
1441
1442impl<'a> ChainRowIterator<'a> {
1443    fn new(prefix: Vec<Row>, tail: Box<dyn RowIterator + 'a>, schema: Vec<ColumnMetadata>) -> Self {
1444        Self {
1445            prefix: prefix.into_iter(),
1446            tail,
1447            schema,
1448        }
1449    }
1450}
1451
1452impl RowIterator for ChainRowIterator<'_> {
1453    fn next_row(&mut self) -> Option<Result<Row>> {
1454        if let Some(row) = self.prefix.next() {
1455            return Some(Ok(row));
1456        }
1457        self.tail.next_row()
1458    }
1459
1460    fn schema(&self) -> &[ColumnMetadata] {
1461        &self.schema
1462    }
1463}
1464
1465fn estimate_row_bytes(row: &Row) -> u64 {
1466    row.values.iter().map(ByteSized::estimated_bytes).sum()
1467}
1468
1469fn split_contiguous_partitions(rows: Vec<Row>, parallelism: usize) -> Vec<Vec<Row>> {
1470    let requested = parallelism.max(1);
1471    if rows.is_empty() {
1472        return (0..requested).map(|_| Vec::new()).collect();
1473    }
1474    let partitions = requested.min(rows.len());
1475    let total = rows.len();
1476    let mut tail = rows;
1477    let mut output = Vec::with_capacity(partitions);
1478    for partition in (0..partitions).rev() {
1479        let start = partition * total / partitions;
1480        output.push(tail.split_off(start));
1481    }
1482    output.reverse();
1483    output
1484}
1485
1486#[allow(clippy::too_many_arguments)]
1487fn execute_partial_partition(
1488    partition_index: usize,
1489    rows: Vec<Row>,
1490    input_schema: Vec<ColumnMetadata>,
1491    group_keys: Vec<TypedExpr>,
1492    aggregates: Vec<AggregateExpr>,
1493    partial_schema: Vec<ColumnMetadata>,
1494    group_limit: usize,
1495    shared_group_counter: Arc<AtomicUsize>,
1496    shared_memory_counter: Arc<AtomicU64>,
1497    memory_limit: Option<u64>,
1498) -> Result<(usize, Vec<Row>)> {
1499    let input = VecIterator::new(rows, input_schema);
1500    let mut iter = AggregateIterator::new(
1501        Box::new(input),
1502        group_keys,
1503        aggregates,
1504        None,
1505        partial_schema,
1506    )
1507    .with_mode(AggregateMode::Partial)
1508    .with_group_limit(group_limit)
1509    .with_shared_group_counter(Some(shared_group_counter));
1510    let rows = collect_iterator_rows(&mut iter)?;
1511    for row in &rows {
1512        let used = shared_memory_counter
1513            .fetch_add(estimate_row_bytes(row), std::sync::atomic::Ordering::SeqCst)
1514            .saturating_add(estimate_row_bytes(row));
1515        if let Some(limit) = memory_limit
1516            && used > limit
1517        {
1518            return Err(ExecutorError::ResourceExhausted {
1519                message: format!(
1520                    "parallel aggregate memory limit exceeded: {used} bytes (limit {limit})"
1521                ),
1522            });
1523        }
1524    }
1525    Ok((partition_index, rows))
1526}
1527
1528fn recv_partition_results(
1529    receiver: std::sync::mpsc::Receiver<Result<(usize, Vec<Row>)>>,
1530    expected: usize,
1531) -> Result<Vec<(usize, Vec<Row>)>> {
1532    let mut outputs = Vec::with_capacity(expected);
1533    for _ in 0..expected {
1534        let result = receiver
1535            .recv()
1536            .map_err(|err| ExecutorError::InvalidOperation {
1537                operation: "parallel aggregate".into(),
1538                reason: format!("partition worker failed to report result: {err}"),
1539            })?;
1540        outputs.push(result?);
1541    }
1542    outputs.sort_by_key(|(idx, _)| *idx);
1543    Ok(outputs)
1544}
1545
1546#[cfg(feature = "tokio")]
1547#[allow(clippy::too_many_arguments)]
1548fn run_partial_partitions(
1549    partitions: Vec<Vec<Row>>,
1550    input_schema: Vec<ColumnMetadata>,
1551    group_keys: Vec<TypedExpr>,
1552    aggregates: Vec<AggregateExpr>,
1553    partial_schema: Vec<ColumnMetadata>,
1554    group_limit: usize,
1555    shared_group_counter: Arc<AtomicUsize>,
1556    shared_memory_counter: Arc<AtomicU64>,
1557    memory_limit: Option<u64>,
1558) -> Result<Vec<(usize, Vec<Row>)>> {
1559    if let Ok(handle) = tokio::runtime::Handle::try_current() {
1560        let expected = partitions.len();
1561        let (sender, receiver) = std::sync::mpsc::channel();
1562        for (partition_index, rows) in partitions.into_iter().enumerate() {
1563            let sender = sender.clone();
1564            let input_schema = input_schema.clone();
1565            let group_keys = group_keys.clone();
1566            let aggregates = aggregates.clone();
1567            let partial_schema = partial_schema.clone();
1568            let shared_group_counter = Arc::clone(&shared_group_counter);
1569            let shared_memory_counter = Arc::clone(&shared_memory_counter);
1570            handle.spawn_blocking(move || {
1571                let result = execute_partial_partition(
1572                    partition_index,
1573                    rows,
1574                    input_schema,
1575                    group_keys,
1576                    aggregates,
1577                    partial_schema,
1578                    group_limit,
1579                    shared_group_counter,
1580                    shared_memory_counter,
1581                    memory_limit,
1582                );
1583                let _ = sender.send(result);
1584            });
1585        }
1586        drop(sender);
1587        return recv_partition_results(receiver, expected);
1588    }
1589
1590    run_partial_partitions_on_threads(
1591        partitions,
1592        input_schema,
1593        group_keys,
1594        aggregates,
1595        partial_schema,
1596        group_limit,
1597        shared_group_counter,
1598        shared_memory_counter,
1599        memory_limit,
1600    )
1601}
1602
1603#[cfg(not(feature = "tokio"))]
1604#[allow(clippy::too_many_arguments)]
1605fn run_partial_partitions(
1606    partitions: Vec<Vec<Row>>,
1607    input_schema: Vec<ColumnMetadata>,
1608    group_keys: Vec<TypedExpr>,
1609    aggregates: Vec<AggregateExpr>,
1610    partial_schema: Vec<ColumnMetadata>,
1611    group_limit: usize,
1612    shared_group_counter: Arc<AtomicUsize>,
1613    shared_memory_counter: Arc<AtomicU64>,
1614    memory_limit: Option<u64>,
1615) -> Result<Vec<(usize, Vec<Row>)>> {
1616    run_partial_partitions_on_threads(
1617        partitions,
1618        input_schema,
1619        group_keys,
1620        aggregates,
1621        partial_schema,
1622        group_limit,
1623        shared_group_counter,
1624        shared_memory_counter,
1625        memory_limit,
1626    )
1627}
1628
1629#[allow(clippy::too_many_arguments)]
1630fn run_partial_partitions_on_threads(
1631    partitions: Vec<Vec<Row>>,
1632    input_schema: Vec<ColumnMetadata>,
1633    group_keys: Vec<TypedExpr>,
1634    aggregates: Vec<AggregateExpr>,
1635    partial_schema: Vec<ColumnMetadata>,
1636    group_limit: usize,
1637    shared_group_counter: Arc<AtomicUsize>,
1638    shared_memory_counter: Arc<AtomicU64>,
1639    memory_limit: Option<u64>,
1640) -> Result<Vec<(usize, Vec<Row>)>> {
1641    let expected = partitions.len();
1642    let (sender, receiver) = std::sync::mpsc::channel();
1643    std::thread::scope(|scope| {
1644        for (partition_index, rows) in partitions.into_iter().enumerate() {
1645            let sender = sender.clone();
1646            let input_schema = input_schema.clone();
1647            let group_keys = group_keys.clone();
1648            let aggregates = aggregates.clone();
1649            let partial_schema = partial_schema.clone();
1650            let shared_group_counter = Arc::clone(&shared_group_counter);
1651            let shared_memory_counter = Arc::clone(&shared_memory_counter);
1652            scope.spawn(move || {
1653                let result = execute_partial_partition(
1654                    partition_index,
1655                    rows,
1656                    input_schema,
1657                    group_keys,
1658                    aggregates,
1659                    partial_schema,
1660                    group_limit,
1661                    shared_group_counter,
1662                    shared_memory_counter,
1663                    memory_limit,
1664                );
1665                let _ = sender.send(result);
1666            });
1667        }
1668    });
1669    drop(sender);
1670    recv_partition_results(receiver, expected)
1671}
1672
1673/// Execute a deterministic single-process parallel partial-to-final aggregate.
1674pub fn execute_parallel_aggregate_rows<'a>(
1675    input: Box<dyn RowIterator + 'a>,
1676    group_keys: Vec<TypedExpr>,
1677    aggregates: Vec<AggregateExpr>,
1678    having: Option<TypedExpr>,
1679    final_schema: Vec<ColumnMetadata>,
1680    parallelism: usize,
1681) -> Result<Vec<Row>> {
1682    execute_parallel_aggregate_rows_with_policy(
1683        input,
1684        group_keys,
1685        aggregates,
1686        having,
1687        final_schema,
1688        parallelism,
1689        None,
1690        DEFAULT_GROUP_LIMIT,
1691    )
1692}
1693
1694/// Execute a deterministic parallel aggregate with memory fallback controls.
1695#[allow(clippy::too_many_arguments)]
1696pub fn execute_parallel_aggregate_rows_with_policy<'a>(
1697    mut input: Box<dyn RowIterator + 'a>,
1698    group_keys: Vec<TypedExpr>,
1699    aggregates: Vec<AggregateExpr>,
1700    having: Option<TypedExpr>,
1701    final_schema: Vec<ColumnMetadata>,
1702    parallelism: usize,
1703    memory: Option<MemoryPolicy>,
1704    group_limit: usize,
1705) -> Result<Vec<Row>> {
1706    if parallelism <= 1 {
1707        return execute_single_aggregate_rows(
1708            input,
1709            group_keys,
1710            aggregates,
1711            having,
1712            final_schema,
1713            memory,
1714            group_limit,
1715        );
1716    }
1717
1718    let input_schema = input.schema().to_vec();
1719    let mut input_rows = Vec::new();
1720    let mut materialized_bytes = 0u64;
1721    let materialize_threshold = memory
1722        .as_ref()
1723        .and_then(MemoryPolicy::limit_bytes)
1724        .map(|limit| (limit / 2).max(1));
1725
1726    while let Some(result) = input.next_row() {
1727        let row = result?;
1728        let row_bytes = materialize_threshold.map(|_| estimate_row_bytes(&row));
1729        if let (Some(threshold), Some(row_bytes)) = (materialize_threshold, row_bytes)
1730            && materialized_bytes.saturating_add(row_bytes) > threshold
1731        {
1732            input_rows.push(row);
1733            let chained = ChainRowIterator::new(input_rows, input, input_schema.clone());
1734            return execute_single_aggregate_rows(
1735                Box::new(chained),
1736                group_keys,
1737                aggregates,
1738                having,
1739                final_schema,
1740                memory,
1741                group_limit,
1742            );
1743        }
1744        if let Some(row_bytes) = row_bytes {
1745            materialized_bytes = materialized_bytes.saturating_add(row_bytes);
1746        }
1747        input_rows.push(row);
1748    }
1749
1750    let fallback_rows = if memory.is_some() || group_limit < input_rows.len() {
1751        Some(input_rows.clone())
1752    } else {
1753        None
1754    };
1755    let result = execute_parallel_aggregate_rows_from_materialized(
1756        input_rows,
1757        input_schema.clone(),
1758        group_keys.clone(),
1759        aggregates.clone(),
1760        having.clone(),
1761        final_schema.clone(),
1762        parallelism,
1763        group_limit,
1764        materialized_bytes,
1765        memory.as_ref().and_then(MemoryPolicy::limit_bytes),
1766    );
1767    match result {
1768        Ok(rows) => Ok(rows),
1769        Err(ExecutorError::ResourceExhausted { .. }) => {
1770            if let Some(fallback_rows) = fallback_rows {
1771                execute_single_aggregate_rows(
1772                    Box::new(VecIterator::new(fallback_rows, input_schema)),
1773                    group_keys,
1774                    aggregates,
1775                    having,
1776                    final_schema,
1777                    memory,
1778                    group_limit,
1779                )
1780            } else {
1781                Err(ExecutorError::ResourceExhausted {
1782                    message: format!(
1783                        "parallel aggregate exceeded group limit {group_limit}; no fallback rows retained"
1784                    ),
1785                })
1786            }
1787        }
1788        Err(err) => Err(err),
1789    }
1790}
1791
1792#[allow(clippy::too_many_arguments)]
1793fn execute_parallel_aggregate_rows_from_materialized(
1794    input_rows: Vec<Row>,
1795    input_schema: Vec<ColumnMetadata>,
1796    group_keys: Vec<TypedExpr>,
1797    aggregates: Vec<AggregateExpr>,
1798    having: Option<TypedExpr>,
1799    final_schema: Vec<ColumnMetadata>,
1800    parallelism: usize,
1801    group_limit: usize,
1802    materialized_bytes: u64,
1803    memory_limit: Option<u64>,
1804) -> Result<Vec<Row>> {
1805    let partial_schema = build_partial_aggregate_schema(&group_keys, &aggregates);
1806    let partitions = split_contiguous_partitions(input_rows, parallelism);
1807    let shared_group_counter = Arc::new(AtomicUsize::new(0));
1808    let shared_memory_counter = Arc::new(AtomicU64::new(materialized_bytes));
1809    let partial_results = run_partial_partitions(
1810        partitions,
1811        input_schema,
1812        group_keys.clone(),
1813        aggregates.clone(),
1814        partial_schema.clone(),
1815        group_limit,
1816        shared_group_counter,
1817        shared_memory_counter,
1818        memory_limit,
1819    )?;
1820
1821    let partial_rows = partial_results
1822        .into_iter()
1823        .flat_map(|(_, rows)| rows)
1824        .collect::<Vec<_>>();
1825    let final_input = VecIterator::new(partial_rows, partial_schema);
1826    let mut final_iter = AggregateIterator::new(
1827        Box::new(final_input),
1828        group_keys,
1829        aggregates,
1830        having,
1831        final_schema,
1832    )
1833    .with_mode(AggregateMode::Final)
1834    .with_group_limit(group_limit);
1835    collect_iterator_rows(&mut final_iter)
1836}
1837
1838fn execute_single_aggregate_rows<'a>(
1839    input: Box<dyn RowIterator + 'a>,
1840    group_keys: Vec<TypedExpr>,
1841    aggregates: Vec<AggregateExpr>,
1842    having: Option<TypedExpr>,
1843    final_schema: Vec<ColumnMetadata>,
1844    memory: Option<MemoryPolicy>,
1845    group_limit: usize,
1846) -> Result<Vec<Row>> {
1847    let mut iter = AggregateIterator::new(input, group_keys, aggregates, having, final_schema)
1848        .with_group_limit(group_limit)
1849        .with_memory_policy(memory);
1850    collect_iterator_rows(&mut iter)
1851}
1852
1853#[cfg(test)]
1854mod tests {
1855    use super::*;
1856    use crate::ast::span::Span;
1857    use crate::executor::memory::SpillPolicy;
1858    use crate::planner::typed_expr::TypedExprKind;
1859
1860    fn apply_values(acc: &mut dyn Accumulator, values: &[Option<SqlValue>]) {
1861        for value in values {
1862            acc.update(value.clone()).unwrap();
1863        }
1864    }
1865
1866    fn single_result(
1867        make_accumulator: impl Fn() -> Box<dyn Accumulator>,
1868        partitions: &[Vec<Option<SqlValue>>],
1869    ) -> SqlValue {
1870        let mut acc = make_accumulator();
1871        for partition in partitions {
1872            apply_values(acc.as_mut(), partition);
1873        }
1874        acc.finalize().unwrap()
1875    }
1876
1877    fn merged_result(
1878        make_partial: impl Fn() -> Box<dyn Accumulator>,
1879        make_final: impl Fn() -> Box<dyn Accumulator>,
1880        partitions: &[Vec<Option<SqlValue>>],
1881        merge_order: &[usize],
1882    ) -> SqlValue {
1883        let states = partitions
1884            .iter()
1885            .map(|partition| {
1886                let mut acc = make_partial();
1887                apply_values(acc.as_mut(), partition);
1888                acc.state().unwrap()
1889            })
1890            .collect::<Vec<_>>();
1891
1892        let mut final_acc = make_final();
1893        for idx in merge_order {
1894            final_acc.merge(&states[*idx]).unwrap();
1895        }
1896        final_acc.finalize().unwrap()
1897    }
1898
1899    fn assert_single_equals_merged(
1900        make_accumulator: impl Fn() -> Box<dyn Accumulator> + Copy,
1901        partitions: Vec<Vec<Option<SqlValue>>>,
1902    ) {
1903        let merge_order = (0..partitions.len()).collect::<Vec<_>>();
1904        let single = single_result(make_accumulator, &partitions);
1905        let merged = merged_result(
1906            make_accumulator,
1907            make_accumulator,
1908            &partitions,
1909            &merge_order,
1910        );
1911        assert_eq!(single, merged);
1912    }
1913
1914    fn assert_merge_order_invariant(
1915        make_accumulator: impl Fn() -> Box<dyn Accumulator> + Copy,
1916        partitions: Vec<Vec<Option<SqlValue>>>,
1917        merge_orders: &[Vec<usize>],
1918    ) {
1919        let single = single_result(make_accumulator, &partitions);
1920        for order in merge_orders {
1921            let merged = merged_result(make_accumulator, make_accumulator, &partitions, order);
1922            assert_eq!(single, merged, "merge order {order:?}");
1923        }
1924    }
1925
1926    fn column_expr(index: usize, name: &str, resolved_type: ResolvedType) -> TypedExpr {
1927        TypedExpr {
1928            kind: TypedExprKind::ColumnRef {
1929                table: "t".into(),
1930                column: name.into(),
1931                column_index: index,
1932            },
1933            resolved_type,
1934            span: Span::default(),
1935        }
1936    }
1937
1938    fn sample_aggregate_schema() -> Vec<ColumnMetadata> {
1939        vec![
1940            ColumnMetadata::new("category", ResolvedType::Text),
1941            ColumnMetadata::new("price", ResolvedType::Double),
1942            ColumnMetadata::new("label", ResolvedType::Text),
1943        ]
1944    }
1945
1946    fn sample_aggregate_rows() -> Vec<Row> {
1947        vec![
1948            Row::new(
1949                0,
1950                vec![
1951                    SqlValue::Text("book".into()),
1952                    SqlValue::Double(10.0),
1953                    SqlValue::Text("a".into()),
1954                ],
1955            ),
1956            Row::new(
1957                1,
1958                vec![
1959                    SqlValue::Text("book".into()),
1960                    SqlValue::Double(15.0),
1961                    SqlValue::Text("b".into()),
1962                ],
1963            ),
1964            Row::new(
1965                2,
1966                vec![
1967                    SqlValue::Text("game".into()),
1968                    SqlValue::Double(20.0),
1969                    SqlValue::Text("c".into()),
1970                ],
1971            ),
1972            Row::new(
1973                3,
1974                vec![
1975                    SqlValue::Text("book".into()),
1976                    SqlValue::Null,
1977                    SqlValue::Text("a".into()),
1978                ],
1979            ),
1980            Row::new(
1981                4,
1982                vec![
1983                    SqlValue::Text("toy".into()),
1984                    SqlValue::Double(3.0),
1985                    SqlValue::Text("d".into()),
1986                ],
1987            ),
1988        ]
1989    }
1990
1991    fn sample_aggregates() -> Vec<AggregateExpr> {
1992        let price = column_expr(1, "price", ResolvedType::Double);
1993        let label = column_expr(2, "label", ResolvedType::Text);
1994        vec![
1995            AggregateExpr::count_star(),
1996            AggregateExpr::sum(price.clone()),
1997            AggregateExpr::avg(price),
1998            AggregateExpr {
1999                function: AggregateFunction::GroupConcat {
2000                    separator: Some("|".into()),
2001                },
2002                arg: Some(label),
2003                distinct: false,
2004                result_type: ResolvedType::Text,
2005            },
2006        ]
2007    }
2008
2009    fn collect_single_aggregate(
2010        group_keys: Vec<TypedExpr>,
2011        aggregates: Vec<AggregateExpr>,
2012    ) -> Vec<Vec<SqlValue>> {
2013        let input = VecIterator::new(sample_aggregate_rows(), sample_aggregate_schema());
2014        let schema = build_aggregate_schema(&group_keys, &aggregates);
2015        let mut iter =
2016            AggregateIterator::new(Box::new(input), group_keys, aggregates, None, schema);
2017        collect_iterator_rows(&mut iter)
2018            .unwrap()
2019            .into_iter()
2020            .map(|row| row.values)
2021            .collect()
2022    }
2023
2024    fn collect_parallel_aggregate(
2025        group_keys: Vec<TypedExpr>,
2026        aggregates: Vec<AggregateExpr>,
2027        parallelism: usize,
2028    ) -> Vec<Vec<SqlValue>> {
2029        let input = VecIterator::new(sample_aggregate_rows(), sample_aggregate_schema());
2030        let schema = build_aggregate_schema(&group_keys, &aggregates);
2031        execute_parallel_aggregate_rows(
2032            Box::new(input),
2033            group_keys,
2034            aggregates,
2035            None,
2036            schema,
2037            parallelism,
2038        )
2039        .unwrap()
2040        .into_iter()
2041        .map(|row| row.values)
2042        .collect()
2043    }
2044
2045    fn sort_rows(mut rows: Vec<Vec<SqlValue>>) -> Vec<Vec<SqlValue>> {
2046        rows.sort_by(|left, right| format!("{left:?}").cmp(&format!("{right:?}")));
2047        rows
2048    }
2049
2050    #[test]
2051    fn partial_schema_uses_group_keys_and_state_columns() {
2052        let category = column_expr(0, "category", ResolvedType::Text);
2053        let price = column_expr(1, "price", ResolvedType::Double);
2054        let aggregates = vec![AggregateExpr::count_star(), AggregateExpr::avg(price)];
2055
2056        let schema = build_partial_aggregate_schema(&[category], &aggregates);
2057        let names = schema
2058            .iter()
2059            .map(|column| column.name.as_str())
2060            .collect::<Vec<_>>();
2061        assert_eq!(
2062            names,
2063            vec![
2064                "category",
2065                "__agg0_state0",
2066                "__agg1_state0",
2067                "__agg1_state1"
2068            ]
2069        );
2070        assert_eq!(schema[1].data_type, ResolvedType::BigInt);
2071        assert_eq!(schema[2].data_type, ResolvedType::Double);
2072        assert_eq!(schema[3].data_type, ResolvedType::BigInt);
2073    }
2074
2075    #[test]
2076    fn parallel_aggregate_matches_single_with_group_by() {
2077        let group_keys = vec![column_expr(0, "category", ResolvedType::Text)];
2078        let aggregates = sample_aggregates();
2079
2080        let single = sort_rows(collect_single_aggregate(
2081            group_keys.clone(),
2082            aggregates.clone(),
2083        ));
2084        let parallel = sort_rows(collect_parallel_aggregate(group_keys, aggregates, 3));
2085
2086        assert_eq!(parallel, single);
2087    }
2088
2089    #[test]
2090    fn parallel_aggregate_matches_single_without_group_by() {
2091        let aggregates = sample_aggregates();
2092
2093        let single = collect_single_aggregate(Vec::new(), aggregates.clone());
2094        let parallel = collect_parallel_aggregate(Vec::new(), aggregates, 4);
2095
2096        assert_eq!(parallel, single);
2097    }
2098
2099    #[test]
2100    fn distinct_aggregates_force_single_parallel_mode() {
2101        let price = column_expr(1, "price", ResolvedType::Double);
2102        let aggregates = vec![AggregateExpr {
2103            distinct: true,
2104            ..AggregateExpr::sum(price)
2105        }];
2106
2107        assert!(should_use_single_for_parallel(4, &aggregates));
2108        assert!(should_use_single_for_parallel(1, &sample_aggregates()));
2109        assert!(!should_use_single_for_parallel(2, &sample_aggregates()));
2110    }
2111
2112    #[test]
2113    fn parallel_group_counter_exhaustion_falls_back_to_single() {
2114        let schema = vec![ColumnMetadata::new("category", ResolvedType::Text)];
2115        let rows = vec![
2116            Row::new(0, vec![SqlValue::Text("a".into())]),
2117            Row::new(1, vec![SqlValue::Text("b".into())]),
2118            Row::new(2, vec![SqlValue::Text("a".into())]),
2119            Row::new(3, vec![SqlValue::Text("b".into())]),
2120        ];
2121        let group_keys = vec![column_expr(0, "category", ResolvedType::Text)];
2122        let aggregates = vec![AggregateExpr::count_star()];
2123        let final_schema = build_aggregate_schema(&group_keys, &aggregates);
2124
2125        let single_values = {
2126            let input = VecIterator::new(rows.clone(), schema.clone());
2127            execute_single_aggregate_rows(
2128                Box::new(input),
2129                group_keys.clone(),
2130                aggregates.clone(),
2131                None,
2132                final_schema.clone(),
2133                None,
2134                2,
2135            )
2136            .unwrap()
2137            .into_iter()
2138            .map(|row| row.values)
2139            .collect::<Vec<_>>()
2140        };
2141
2142        let parallel_values = execute_parallel_aggregate_rows_with_policy(
2143            Box::new(VecIterator::new(rows, schema)),
2144            group_keys,
2145            aggregates,
2146            None,
2147            final_schema,
2148            2,
2149            None,
2150            2,
2151        )
2152        .unwrap()
2153        .into_iter()
2154        .map(|row| row.values)
2155        .collect::<Vec<_>>();
2156
2157        assert_eq!(sort_rows(parallel_values), sort_rows(single_values));
2158    }
2159
2160    #[test]
2161    fn materialize_limit_exhaustion_falls_back_to_streaming_single() {
2162        let schema = vec![ColumnMetadata::new("payload", ResolvedType::Text)];
2163        let rows = (0..4)
2164            .map(|idx| Row::new(idx, vec![SqlValue::Text("x".repeat(40))]))
2165            .collect::<Vec<_>>();
2166        let aggregates = vec![AggregateExpr::count_star()];
2167        let final_schema = build_aggregate_schema(&[], &aggregates);
2168        let policy = MemoryPolicy::new(Some(100), SpillPolicy::FailFast);
2169
2170        let result = execute_parallel_aggregate_rows_with_policy(
2171            Box::new(VecIterator::new(rows, schema)),
2172            Vec::new(),
2173            aggregates,
2174            None,
2175            final_schema,
2176            4,
2177            Some(policy),
2178            DEFAULT_GROUP_LIMIT,
2179        )
2180        .unwrap();
2181
2182        assert_eq!(result.len(), 1);
2183        assert_eq!(result[0].values, vec![SqlValue::BigInt(4)]);
2184    }
2185
2186    #[test]
2187    fn streaming_aggregate_respects_distinct_accumulators() {
2188        let schema = sample_aggregate_schema();
2189        let rows = vec![
2190            Row::new(
2191                0,
2192                vec![
2193                    SqlValue::Text("book".into()),
2194                    SqlValue::Double(10.0),
2195                    SqlValue::Text("a".into()),
2196                ],
2197            ),
2198            Row::new(
2199                1,
2200                vec![
2201                    SqlValue::Text("book".into()),
2202                    SqlValue::Double(10.0),
2203                    SqlValue::Text("a".into()),
2204                ],
2205            ),
2206            Row::new(
2207                2,
2208                vec![
2209                    SqlValue::Text("book".into()),
2210                    SqlValue::Double(15.0),
2211                    SqlValue::Text("b".into()),
2212                ],
2213            ),
2214        ];
2215        let group_keys = vec![column_expr(0, "category", ResolvedType::Text)];
2216        let price = column_expr(1, "price", ResolvedType::Double);
2217        let label = column_expr(2, "label", ResolvedType::Text);
2218        let aggregates = vec![
2219            AggregateExpr {
2220                distinct: true,
2221                ..AggregateExpr::sum(price)
2222            },
2223            AggregateExpr {
2224                function: AggregateFunction::GroupConcat {
2225                    separator: Some("|".into()),
2226                },
2227                arg: Some(label),
2228                distinct: true,
2229                result_type: ResolvedType::Text,
2230            },
2231        ];
2232        let output_schema = build_aggregate_schema(&group_keys, &aggregates);
2233        let input = VecIterator::new(rows, schema);
2234        let mut iter = StreamingAggregateIterator::new(
2235            Box::new(input),
2236            group_keys,
2237            aggregates,
2238            None,
2239            output_schema,
2240        );
2241        let rows = collect_iterator_rows(&mut iter).unwrap();
2242
2243        assert_eq!(rows.len(), 1);
2244        assert_eq!(rows[0].values[1], SqlValue::Double(25.0));
2245        assert_eq!(rows[0].values[2], SqlValue::Text("a|b".into()));
2246    }
2247
2248    #[test]
2249    fn partial_state_matches_single_for_count_sum_total_avg_min_max() {
2250        assert_single_equals_merged(
2251            || Box::new(CountAccumulator::new(false)),
2252            vec![
2253                vec![
2254                    Some(SqlValue::Integer(1)),
2255                    Some(SqlValue::BigInt(2)),
2256                    Some(SqlValue::Text("x".into())),
2257                    Some(SqlValue::Null),
2258                ],
2259                vec![Some(SqlValue::Integer(3))],
2260            ],
2261        );
2262        assert_single_equals_merged(
2263            || Box::new(SumAccumulator::new()),
2264            vec![
2265                vec![
2266                    Some(SqlValue::Integer(1)),
2267                    Some(SqlValue::BigInt(2)),
2268                    Some(SqlValue::Float(3.5)),
2269                ],
2270                vec![Some(SqlValue::Double(4.5)), Some(SqlValue::Null)],
2271            ],
2272        );
2273        assert_single_equals_merged(
2274            || Box::new(TotalAccumulator::new()),
2275            vec![
2276                vec![Some(SqlValue::Integer(1)), Some(SqlValue::Null)],
2277                vec![Some(SqlValue::Double(2.5))],
2278            ],
2279        );
2280        assert_single_equals_merged(
2281            || Box::new(AvgAccumulator::new()),
2282            vec![
2283                vec![Some(SqlValue::Integer(2)), Some(SqlValue::Double(4.0))],
2284                vec![Some(SqlValue::Null), Some(SqlValue::Double(6.0))],
2285            ],
2286        );
2287        assert_single_equals_merged(
2288            || Box::new(MinMaxAccumulator::new(true)),
2289            vec![
2290                vec![Some(SqlValue::Integer(3)), Some(SqlValue::Integer(1))],
2291                vec![Some(SqlValue::Integer(2)), Some(SqlValue::Null)],
2292            ],
2293        );
2294        assert_single_equals_merged(
2295            || Box::new(MinMaxAccumulator::new(false)),
2296            vec![
2297                vec![
2298                    Some(SqlValue::Text("b".into())),
2299                    Some(SqlValue::Text("a".into())),
2300                ],
2301                vec![Some(SqlValue::Text("c".into())), Some(SqlValue::Null)],
2302            ],
2303        );
2304    }
2305
2306    #[test]
2307    fn partial_state_matches_single_for_ordered_string_aggregates() {
2308        assert_single_equals_merged(
2309            || Box::new(GroupConcatAccumulator::new("|".into())),
2310            vec![
2311                vec![Some(SqlValue::Text("a".into())), Some(SqlValue::Null)],
2312                vec![
2313                    Some(SqlValue::Text("b".into())),
2314                    Some(SqlValue::Text("c".into())),
2315                ],
2316            ],
2317        );
2318        assert_single_equals_merged(
2319            || Box::new(StringAggAccumulator::new("::".into())),
2320            vec![
2321                vec![Some(SqlValue::Text("a".into()))],
2322                vec![Some(SqlValue::Text("b".into())), Some(SqlValue::Null)],
2323            ],
2324        );
2325    }
2326
2327    #[test]
2328    fn partial_state_handles_empty_all_null_single_and_mixed_boundaries() {
2329        assert_single_equals_merged(
2330            || Box::new(CountAccumulator::new(false)),
2331            vec![vec![], vec![]],
2332        );
2333        assert_single_equals_merged(|| Box::new(SumAccumulator::new()), vec![vec![], vec![]]);
2334        assert_single_equals_merged(|| Box::new(TotalAccumulator::new()), vec![vec![], vec![]]);
2335        assert_single_equals_merged(|| Box::new(AvgAccumulator::new()), vec![vec![], vec![]]);
2336        assert_single_equals_merged(
2337            || Box::new(MinMaxAccumulator::new(true)),
2338            vec![vec![], vec![]],
2339        );
2340        assert_single_equals_merged(
2341            || Box::new(MinMaxAccumulator::new(false)),
2342            vec![vec![], vec![]],
2343        );
2344        assert_single_equals_merged(
2345            || Box::new(GroupConcatAccumulator::new(",".into())),
2346            vec![vec![], vec![]],
2347        );
2348        assert_single_equals_merged(
2349            || Box::new(StringAggAccumulator::new(",".into())),
2350            vec![vec![], vec![]],
2351        );
2352        assert_single_equals_merged(
2353            || Box::new(CountAccumulator::new(false)),
2354            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2355        );
2356        assert_single_equals_merged(
2357            || Box::new(SumAccumulator::new()),
2358            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2359        );
2360        assert_single_equals_merged(
2361            || Box::new(TotalAccumulator::new()),
2362            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2363        );
2364        assert_single_equals_merged(
2365            || Box::new(AvgAccumulator::new()),
2366            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2367        );
2368        assert_single_equals_merged(
2369            || Box::new(MinMaxAccumulator::new(true)),
2370            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2371        );
2372        assert_single_equals_merged(
2373            || Box::new(MinMaxAccumulator::new(false)),
2374            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2375        );
2376        assert_single_equals_merged(
2377            || Box::new(GroupConcatAccumulator::new(",".into())),
2378            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2379        );
2380        assert_single_equals_merged(
2381            || Box::new(StringAggAccumulator::new(",".into())),
2382            vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
2383        );
2384        assert_single_equals_merged(
2385            || Box::new(SumAccumulator::new()),
2386            vec![vec![Some(SqlValue::Integer(7))]],
2387        );
2388        assert_single_equals_merged(
2389            || Box::new(AvgAccumulator::new()),
2390            vec![
2391                vec![Some(SqlValue::Null)],
2392                vec![Some(SqlValue::Double(8.0))],
2393            ],
2394        );
2395    }
2396
2397    #[test]
2398    fn commutative_accumulators_are_merge_order_invariant() {
2399        let orders = vec![
2400            vec![1, 3, 0, 2],
2401            vec![3, 2, 1, 0],
2402            vec![0, 1, 2, 3],
2403            vec![2, 0, 3, 1],
2404        ];
2405        assert_merge_order_invariant(
2406            || Box::new(CountAccumulator::new(false)),
2407            vec![vec![Some(SqlValue::Integer(1))]],
2408            &[vec![0]],
2409        );
2410        assert_merge_order_invariant(
2411            || Box::new(SumAccumulator::new()),
2412            vec![
2413                vec![Some(SqlValue::Integer(1))],
2414                vec![Some(SqlValue::Integer(2))],
2415            ],
2416            &[vec![0, 1], vec![1, 0]],
2417        );
2418        assert_merge_order_invariant(
2419            || Box::new(AvgAccumulator::new()),
2420            vec![
2421                vec![Some(SqlValue::Integer(1))],
2422                vec![Some(SqlValue::Integer(2))],
2423                vec![Some(SqlValue::Integer(3))],
2424            ],
2425            &[vec![0, 1, 2], vec![2, 1, 0]],
2426        );
2427        let numeric_partitions = vec![
2428            vec![Some(SqlValue::Integer(1)), Some(SqlValue::Null)],
2429            vec![Some(SqlValue::BigInt(2))],
2430            vec![],
2431            vec![Some(SqlValue::Double(3.0))],
2432        ];
2433        assert_merge_order_invariant(
2434            || Box::new(CountAccumulator::new(false)),
2435            numeric_partitions.clone(),
2436            &orders,
2437        );
2438        assert_merge_order_invariant(
2439            || Box::new(SumAccumulator::new()),
2440            numeric_partitions.clone(),
2441            &orders,
2442        );
2443        assert_merge_order_invariant(
2444            || Box::new(TotalAccumulator::new()),
2445            numeric_partitions.clone(),
2446            &orders,
2447        );
2448        assert_merge_order_invariant(
2449            || Box::new(AvgAccumulator::new()),
2450            numeric_partitions.clone(),
2451            &orders,
2452        );
2453        let integer_partitions = vec![
2454            vec![Some(SqlValue::Integer(3)), Some(SqlValue::Null)],
2455            vec![Some(SqlValue::Integer(1))],
2456            vec![],
2457            vec![Some(SqlValue::Integer(2))],
2458        ];
2459        assert_merge_order_invariant(
2460            || Box::new(MinMaxAccumulator::new(true)),
2461            integer_partitions.clone(),
2462            &orders,
2463        );
2464        assert_merge_order_invariant(
2465            || Box::new(MinMaxAccumulator::new(false)),
2466            integer_partitions,
2467            &orders,
2468        );
2469    }
2470
2471    #[test]
2472    fn avg_partial_state_uses_sum_count_and_never_divides_by_zero_during_merge() {
2473        let empty = {
2474            let acc = AvgAccumulator::new();
2475            acc.state().unwrap()
2476        };
2477        assert_eq!(empty, vec![SqlValue::Double(0.0), SqlValue::BigInt(0)]);
2478
2479        let mut partial = AvgAccumulator::new();
2480        partial.update(Some(SqlValue::Integer(2))).unwrap();
2481        partial.update(Some(SqlValue::Double(4.0))).unwrap();
2482        assert_eq!(
2483            partial.state().unwrap(),
2484            vec![SqlValue::Double(6.0), SqlValue::BigInt(2)]
2485        );
2486
2487        let mut final_acc = AvgAccumulator::new();
2488        final_acc.merge(&empty).unwrap();
2489        assert_eq!(final_acc.finalize().unwrap(), SqlValue::Null);
2490        final_acc.merge(&partial.state().unwrap()).unwrap();
2491        assert_eq!(final_acc.finalize().unwrap(), SqlValue::Double(3.0));
2492    }
2493
2494    #[test]
2495    fn merge_rejects_invalid_state_contracts_without_panicking() {
2496        let mut count = CountAccumulator::new(false);
2497        assert!(count.merge(&[]).is_err());
2498        assert!(count.merge(&[SqlValue::Text("bad".into())]).is_err());
2499
2500        let mut avg = AvgAccumulator::new();
2501        assert!(avg.merge(&[SqlValue::Double(1.0)]).is_err());
2502        assert!(
2503            avg.merge(&[SqlValue::Double(1.0), SqlValue::Text("bad".into())])
2504                .is_err()
2505        );
2506
2507        let mut concat = GroupConcatAccumulator::new("|".into());
2508        assert!(
2509            concat
2510                .merge(&[SqlValue::Text("a".into()), SqlValue::Text(",".into())])
2511                .is_err()
2512        );
2513    }
2514
2515    #[test]
2516    fn count_accumulator_counts_rows_and_skips_nulls() {
2517        let mut acc = CountAccumulator::new(false);
2518        acc.update(None).unwrap();
2519        acc.update(Some(SqlValue::Null)).unwrap();
2520        acc.update(Some(SqlValue::Integer(1))).unwrap();
2521        assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(2));
2522    }
2523
2524    #[test]
2525    fn count_accumulator_distinct_deduplicates() {
2526        let mut acc = CountAccumulator::new(true);
2527        acc.update(Some(SqlValue::Integer(1))).unwrap();
2528        acc.update(Some(SqlValue::Integer(1))).unwrap();
2529        acc.update(Some(SqlValue::Integer(2))).unwrap();
2530        assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(2));
2531    }
2532
2533    #[test]
2534    fn count_distinct_uses_group_key_equality_boundaries() {
2535        let mut acc = CountAccumulator::new(true);
2536        let nan_a = f64::from_bits(0x7ff8_0000_0000_0001);
2537        let nan_b = f64::from_bits(0x7ff8_0000_0000_0002);
2538        for value in [
2539            SqlValue::Null,
2540            SqlValue::Null,
2541            SqlValue::Integer(1),
2542            SqlValue::Integer(1),
2543            SqlValue::Double(1.0),
2544            SqlValue::Double(-0.0),
2545            SqlValue::Double(0.0),
2546            SqlValue::Double(nan_a),
2547            SqlValue::Double(nan_a),
2548            SqlValue::Double(nan_b),
2549            SqlValue::Text("same".into()),
2550            SqlValue::Text("same".into()),
2551            SqlValue::Blob(vec![1, 2]),
2552            SqlValue::Blob(vec![1, 2]),
2553            SqlValue::Blob(vec![1, 3]),
2554        ] {
2555            acc.update(Some(value)).unwrap();
2556        }
2557        assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(9));
2558    }
2559
2560    #[test]
2561    fn distinct_non_count_accumulators_deduplicate_non_null_values() {
2562        let mut sum = SumAccumulator::with_distinct(true);
2563        for value in [
2564            SqlValue::Integer(1),
2565            SqlValue::Integer(1),
2566            SqlValue::Double(1.0),
2567            SqlValue::Integer(2),
2568            SqlValue::Null,
2569        ] {
2570            sum.update(Some(value)).unwrap();
2571        }
2572        assert_eq!(sum.finalize().unwrap(), SqlValue::Double(4.0));
2573
2574        let mut avg = AvgAccumulator::with_distinct(true);
2575        for value in [
2576            SqlValue::Integer(1),
2577            SqlValue::Integer(1),
2578            SqlValue::Integer(3),
2579            SqlValue::Null,
2580        ] {
2581            avg.update(Some(value)).unwrap();
2582        }
2583        assert_eq!(avg.finalize().unwrap(), SqlValue::Double(2.0));
2584
2585        let mut min = MinMaxAccumulator::with_distinct(true, true);
2586        let mut max = MinMaxAccumulator::with_distinct(false, true);
2587        for value in [
2588            SqlValue::Text("b".into()),
2589            SqlValue::Text("a".into()),
2590            SqlValue::Text("a".into()),
2591            SqlValue::Text("c".into()),
2592        ] {
2593            min.update(Some(value.clone())).unwrap();
2594            max.update(Some(value)).unwrap();
2595        }
2596        assert_eq!(min.finalize().unwrap(), SqlValue::Text("a".into()));
2597        assert_eq!(max.finalize().unwrap(), SqlValue::Text("c".into()));
2598
2599        let mut group_concat = GroupConcatAccumulator::with_distinct("|".into(), true);
2600        let mut string_agg = StringAggAccumulator::with_distinct(";".into(), true);
2601        for value in [
2602            SqlValue::Text("a".into()),
2603            SqlValue::Text("a".into()),
2604            SqlValue::Null,
2605            SqlValue::Text("b".into()),
2606        ] {
2607            group_concat.update(Some(value.clone())).unwrap();
2608            string_agg.update(Some(value)).unwrap();
2609        }
2610        assert_eq!(
2611            group_concat.finalize().unwrap(),
2612            SqlValue::Text("a|b".into())
2613        );
2614        assert_eq!(string_agg.finalize().unwrap(), SqlValue::Text("a;b".into()));
2615    }
2616
2617    #[test]
2618    fn sum_accumulator_aggregates_numeric_values() {
2619        let mut acc = SumAccumulator::new();
2620        acc.update(Some(SqlValue::Integer(2))).unwrap();
2621        acc.update(Some(SqlValue::Double(3.5))).unwrap();
2622        acc.update(Some(SqlValue::Null)).unwrap();
2623        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(5.5));
2624    }
2625
2626    #[test]
2627    fn total_accumulator_returns_zero_for_empty() {
2628        let acc = TotalAccumulator::new();
2629        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(0.0));
2630    }
2631
2632    #[test]
2633    fn total_accumulator_aggregates_numeric_values() {
2634        let mut acc = TotalAccumulator::new();
2635        acc.update(Some(SqlValue::Integer(2))).unwrap();
2636        acc.update(Some(SqlValue::Null)).unwrap();
2637        acc.update(Some(SqlValue::Double(1.5))).unwrap();
2638        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(3.5));
2639    }
2640
2641    #[test]
2642    fn avg_accumulator_handles_empty_and_nulls() {
2643        let mut acc = AvgAccumulator::new();
2644        assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
2645        acc.update(Some(SqlValue::Null)).unwrap();
2646        acc.update(Some(SqlValue::BigInt(4))).unwrap();
2647        acc.update(Some(SqlValue::Integer(2))).unwrap();
2648        assert_eq!(acc.finalize().unwrap(), SqlValue::Double(3.0));
2649    }
2650
2651    #[test]
2652    fn min_max_accumulator_tracks_extremes() {
2653        let mut min_acc = MinMaxAccumulator::new(true);
2654        let mut max_acc = MinMaxAccumulator::new(false);
2655        for value in [3, 1, 2] {
2656            min_acc.update(Some(SqlValue::Integer(value))).unwrap();
2657            max_acc.update(Some(SqlValue::Integer(value))).unwrap();
2658        }
2659        assert_eq!(min_acc.finalize().unwrap(), SqlValue::Integer(1));
2660        assert_eq!(max_acc.finalize().unwrap(), SqlValue::Integer(3));
2661    }
2662
2663    #[test]
2664    fn min_max_accumulator_rejects_type_mismatch() {
2665        let mut acc = MinMaxAccumulator::new(true);
2666        acc.update(Some(SqlValue::Integer(1))).unwrap();
2667        let err = acc.update(Some(SqlValue::Text("bad".into()))).unwrap_err();
2668        match err {
2669            ExecutorError::Evaluation(crate::executor::EvaluationError::TypeMismatch {
2670                ..
2671            }) => {}
2672            other => panic!("unexpected error {:?}", other),
2673        }
2674    }
2675
2676    #[test]
2677    fn group_concat_accumulator_joins_values() {
2678        let mut acc = GroupConcatAccumulator::new("|".into());
2679        acc.update(Some(SqlValue::Text("a".into()))).unwrap();
2680        acc.update(Some(SqlValue::Null)).unwrap();
2681        acc.update(Some(SqlValue::Text("b".into()))).unwrap();
2682        assert_eq!(acc.finalize().unwrap(), SqlValue::Text("a|b".into()));
2683    }
2684
2685    #[test]
2686    fn group_concat_accumulator_empty_returns_null() {
2687        let acc = GroupConcatAccumulator::new(",".into());
2688        assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
2689    }
2690
2691    #[test]
2692    fn string_agg_accumulator_joins_values() {
2693        let mut acc = StringAggAccumulator::new("::".into());
2694        acc.update(Some(SqlValue::Text("a".into()))).unwrap();
2695        acc.update(Some(SqlValue::Null)).unwrap();
2696        acc.update(Some(SqlValue::Text("b".into()))).unwrap();
2697        assert_eq!(acc.finalize().unwrap(), SqlValue::Text("a::b".into()));
2698    }
2699
2700    #[test]
2701    fn string_agg_accumulator_empty_returns_null() {
2702        let acc = StringAggAccumulator::new(",".into());
2703        assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
2704    }
2705
2706    #[test]
2707    fn encode_group_key_is_deterministic() {
2708        let values = vec![
2709            SqlValue::Integer(1),
2710            SqlValue::Text("a".into()),
2711            SqlValue::Null,
2712        ];
2713        let first = encode_group_key(&values).unwrap();
2714        let second = encode_group_key(&values).unwrap();
2715        assert_eq!(first, second);
2716    }
2717}