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