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