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
18pub 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
81pub 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
90pub trait Accumulator: Send {
92 fn update(&mut self, value: Option<SqlValue>) -> Result<()>;
94 fn state(&self) -> Result<Vec<SqlValue>>;
96 fn merge(&mut self, state: &[SqlValue]) -> Result<()>;
98 fn finalize(&self) -> Result<SqlValue>;
100 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#[derive(Debug, Clone)]
183pub struct CountAccumulator {
184 count: usize,
185 distinct_values: Option<HashSet<Vec<u8>>>,
186}
187
188impl CountAccumulator {
189 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#[derive(Debug, Clone)]
253pub struct SumAccumulator {
254 sum: Option<f64>,
255 distinct_values: Option<HashSet<Vec<u8>>>,
256}
257
258impl SumAccumulator {
259 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#[derive(Debug, Clone)]
319pub struct TotalAccumulator {
320 sum: Option<f64>,
321}
322
323impl TotalAccumulator {
324 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#[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 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#[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 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#[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 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#[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 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
742pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
767pub enum AggregateMode {
768 Partial,
770 Final,
772 Single,
774}
775
776struct AggregateGroup {
777 key_values: Vec<SqlValue>,
778 accumulators: Vec<Box<dyn Accumulator>>,
779}
780
781pub 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 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 pub fn with_group_limit(mut self, limit: usize) -> Self {
824 self.group_limit = limit;
825 self
826 }
827
828 pub fn with_mode(mut self, mode: AggregateMode) -> Self {
830 self.mode = mode;
831 self
832 }
833
834 pub fn with_memory_policy(mut self, policy: Option<MemoryPolicy>) -> Self {
836 self.memory_tracker = policy.map(MemoryTracker::new);
837 self
838 }
839
840 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
1072pub 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(¤t_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(¤t_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
1251pub 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
1280pub 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
1304pub 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
1554pub 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#[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}