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 update_values(&mut self, values: &[SqlValue]) -> Result<()> {
97 self.update(values.first().cloned())
98 }
99 fn update_ordered(&mut self, value: Option<SqlValue>, _sort_keys: &[SqlValue]) -> Result<()> {
105 self.update(value)
106 }
107 fn update_ordered_values(&mut self, values: &[SqlValue], sort_keys: &[SqlValue]) -> Result<()> {
108 self.update_ordered(values.first().cloned(), sort_keys)
109 }
110 fn state(&self) -> Result<Vec<SqlValue>>;
112 fn merge(&mut self, state: &[SqlValue]) -> Result<()>;
114 fn finalize(&self) -> Result<SqlValue>;
116 fn clone_box(&self) -> Box<dyn Accumulator>;
118 fn retained_bytes(&self) -> u64 {
123 0
124 }
125}
126
127impl Clone for Box<dyn Accumulator> {
128 fn clone(&self) -> Self {
129 self.clone_box()
130 }
131}
132
133fn invalid_aggregate_state(function: &str, reason: impl Into<String>) -> ExecutorError {
134 ExecutorError::InvalidOperation {
135 operation: function.into(),
136 reason: reason.into(),
137 }
138}
139
140fn expect_state_arity(function: &str, state: &[SqlValue], expected: usize) -> Result<()> {
141 if state.len() == expected {
142 Ok(())
143 } else {
144 Err(invalid_aggregate_state(
145 function,
146 format!("expected {expected} state value(s), got {}", state.len()),
147 ))
148 }
149}
150
151fn state_bigint(function: &str, value: &SqlValue, index: usize) -> Result<i64> {
152 match value {
153 SqlValue::BigInt(v) => Ok(*v),
154 other => Err(invalid_aggregate_state(
155 function,
156 format!(
157 "state value {index} expected BigInt, got {}",
158 other.type_name()
159 ),
160 )),
161 }
162}
163
164fn state_double(function: &str, value: &SqlValue, index: usize) -> Result<f64> {
165 match value {
166 SqlValue::Double(v) => Ok(*v),
167 other => Err(invalid_aggregate_state(
168 function,
169 format!(
170 "state value {index} expected Double, got {}",
171 other.type_name()
172 ),
173 )),
174 }
175}
176
177fn state_text<'a>(function: &str, value: &'a SqlValue, index: usize) -> Result<&'a str> {
178 match value {
179 SqlValue::Text(v) => Ok(v),
180 other => Err(invalid_aggregate_state(
181 function,
182 format!(
183 "state value {index} expected Text, got {}",
184 other.type_name()
185 ),
186 )),
187 }
188}
189
190fn distinct_allows(
191 distinct_values: &mut Option<HashSet<Vec<u8>>>,
192 value: &SqlValue,
193) -> Result<bool> {
194 if value.is_null() {
195 return Ok(false);
196 }
197 let Some(distinct) = distinct_values else {
198 return Ok(true);
199 };
200 let encoded = encode_group_key(std::slice::from_ref(value))?;
201 Ok(distinct.insert(encoded))
202}
203
204fn estimated_distinct_retained_bytes(distinct_values: &Option<HashSet<Vec<u8>>>) -> u64 {
205 let Some(values) = distinct_values else {
206 return 0;
207 };
208 let bucket_bytes = values
211 .capacity()
212 .saturating_mul(std::mem::size_of::<Vec<u8>>() + std::mem::size_of::<u64>() + 1);
213 values.iter().fold(
214 u64::try_from(bucket_bytes).unwrap_or(u64::MAX),
215 |total, value| total.saturating_add(u64::try_from(value.capacity()).unwrap_or(u64::MAX)),
216 )
217}
218
219fn estimated_string_collection_bytes(
220 values: &[String],
221 values_capacity: usize,
222 separator_capacity: usize,
223) -> u64 {
224 let slots = values_capacity.saturating_mul(std::mem::size_of::<String>());
225 values
226 .iter()
227 .fold(u64::try_from(slots).unwrap_or(u64::MAX), |total, value| {
228 total.saturating_add(u64::try_from(value.capacity()).unwrap_or(u64::MAX))
229 })
230 .saturating_add(u64::try_from(separator_capacity).unwrap_or(u64::MAX))
231}
232
233#[derive(Debug, Clone)]
235pub struct CountAccumulator {
236 count: usize,
237 distinct_values: Option<HashSet<Vec<u8>>>,
238}
239
240impl CountAccumulator {
241 pub fn new(distinct: bool) -> Self {
243 Self {
244 count: 0,
245 distinct_values: if distinct { Some(HashSet::new()) } else { None },
246 }
247 }
248}
249
250impl Accumulator for CountAccumulator {
251 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
252 match (&mut self.distinct_values, value) {
253 (Some(distinct), Some(value)) => {
254 if value.is_null() {
255 return Ok(());
256 }
257 let encoded = encode_group_key(std::slice::from_ref(&value))?;
258 if distinct.insert(encoded) {
259 self.count += 1;
260 }
261 }
262 (Some(_), None) => {
263 self.count += 1;
264 }
265 (None, Some(value)) => {
266 if !value.is_null() {
267 self.count += 1;
268 }
269 }
270 (None, None) => {
271 self.count += 1;
272 }
273 }
274 Ok(())
275 }
276
277 fn finalize(&self) -> Result<SqlValue> {
278 Ok(SqlValue::BigInt(self.count as i64))
279 }
280
281 fn state(&self) -> Result<Vec<SqlValue>> {
282 Ok(vec![SqlValue::BigInt(self.count as i64)])
283 }
284
285 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
286 expect_state_arity("count", state, 1)?;
287 let count = state_bigint("count", &state[0], 0)?;
288 if count < 0 {
289 return Err(invalid_aggregate_state(
290 "count",
291 "state count must be non-negative",
292 ));
293 }
294 self.count = self.count.saturating_add(count as usize);
295 Ok(())
296 }
297
298 fn clone_box(&self) -> Box<dyn Accumulator> {
299 Box::new(self.clone())
300 }
301
302 fn retained_bytes(&self) -> u64 {
303 estimated_distinct_retained_bytes(&self.distinct_values)
304 }
305}
306
307#[derive(Debug, Clone)]
309pub struct SumAccumulator {
310 sum: Option<SqlValue>,
311 result_type: ResolvedType,
312 distinct_values: Option<HashSet<Vec<u8>>>,
313}
314
315impl SumAccumulator {
316 pub fn new() -> Self {
318 Self::with_distinct(false)
319 }
320
321 pub fn with_distinct(distinct: bool) -> Self {
322 Self::with_distinct_for_type(distinct, ResolvedType::Double)
323 }
324
325 pub fn with_distinct_for_type(distinct: bool, result_type: ResolvedType) -> Self {
326 Self {
327 sum: None,
328 result_type,
329 distinct_values: if distinct { Some(HashSet::new()) } else { None },
330 }
331 }
332
333 fn add_value(&mut self, value: SqlValue) -> Result<()> {
334 let next = match &self.result_type {
335 ResolvedType::Integer => {
336 let SqlValue::Integer(value) = value else {
337 return sum_type_mismatch("Integer", &value);
338 };
339 let sum = match self.sum.as_ref() {
340 None => value,
341 Some(SqlValue::Integer(current)) => current
342 .checked_add(value)
343 .ok_or(ExecutorError::Evaluation(EvaluationError::Overflow))?,
344 Some(other) => return sum_type_mismatch("Integer", other),
345 };
346 SqlValue::Integer(sum)
347 }
348 ResolvedType::BigInt => {
349 let value = match value {
350 SqlValue::Integer(value) => i64::from(value),
351 SqlValue::BigInt(value) => value,
352 other => return sum_type_mismatch("BigInt", &other),
353 };
354 let sum = match self.sum.as_ref() {
355 None => value,
356 Some(SqlValue::BigInt(current)) => current
357 .checked_add(value)
358 .ok_or(ExecutorError::Evaluation(EvaluationError::Overflow))?,
359 Some(other) => return sum_type_mismatch("BigInt", other),
360 };
361 SqlValue::BigInt(sum)
362 }
363 _ => {
364 let value = numeric_to_f64(&value)?;
365 let sum = match self.sum.as_ref() {
366 None => value,
367 Some(SqlValue::Double(current)) => *current + value,
368 Some(other) => return sum_type_mismatch("Double", other),
369 };
370 SqlValue::Double(sum)
371 }
372 };
373 self.sum = Some(next);
374 Ok(())
375 }
376}
377
378impl Default for SumAccumulator {
379 fn default() -> Self {
380 Self::new()
381 }
382}
383
384impl Accumulator for SumAccumulator {
385 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
386 let Some(value) = value else {
387 return Ok(());
388 };
389 if value.is_null() {
390 return Ok(());
391 }
392 if !distinct_allows(&mut self.distinct_values, &value)? {
393 return Ok(());
394 }
395 self.add_value(value)
396 }
397
398 fn finalize(&self) -> Result<SqlValue> {
399 Ok(self.sum.clone().unwrap_or(SqlValue::Null))
400 }
401
402 fn state(&self) -> Result<Vec<SqlValue>> {
403 Ok(vec![self.sum.clone().unwrap_or(SqlValue::Null)])
404 }
405
406 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
407 expect_state_arity("sum", state, 1)?;
408 if state[0].is_null() {
409 return Ok(());
410 }
411 self.add_value(state[0].clone())
412 }
413
414 fn clone_box(&self) -> Box<dyn Accumulator> {
415 Box::new(self.clone())
416 }
417
418 fn retained_bytes(&self) -> u64 {
419 estimated_distinct_retained_bytes(&self.distinct_values).saturating_add(
420 self.sum
421 .as_ref()
422 .map(ByteSized::estimated_bytes)
423 .unwrap_or(0),
424 )
425 }
426}
427
428#[derive(Debug, Clone)]
430pub struct TotalAccumulator {
431 sum: Option<f64>,
432}
433
434impl TotalAccumulator {
435 pub fn new() -> Self {
437 Self { sum: None }
438 }
439}
440
441impl Default for TotalAccumulator {
442 fn default() -> Self {
443 Self::new()
444 }
445}
446
447impl Accumulator for TotalAccumulator {
448 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
449 let Some(value) = value else {
450 return Ok(());
451 };
452 if value.is_null() {
453 return Ok(());
454 }
455 let numeric = numeric_to_f64(&value)?;
456 self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
457 Ok(())
458 }
459
460 fn finalize(&self) -> Result<SqlValue> {
461 Ok(SqlValue::Double(self.sum.unwrap_or(0.0)))
462 }
463
464 fn state(&self) -> Result<Vec<SqlValue>> {
465 Ok(vec![SqlValue::Double(self.sum.unwrap_or(0.0))])
466 }
467
468 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
469 expect_state_arity("total", state, 1)?;
470 let value = state_double("total", &state[0], 0)?;
471 self.sum = Some(self.sum.unwrap_or(0.0) + value);
472 Ok(())
473 }
474
475 fn clone_box(&self) -> Box<dyn Accumulator> {
476 Box::new(self.clone())
477 }
478}
479
480#[derive(Debug, Clone)]
482pub struct AvgAccumulator {
483 sum: Option<f64>,
484 count: usize,
485 distinct_values: Option<HashSet<Vec<u8>>>,
486}
487
488impl AvgAccumulator {
489 pub fn new() -> Self {
491 Self::with_distinct(false)
492 }
493
494 pub fn with_distinct(distinct: bool) -> Self {
495 Self {
496 sum: None,
497 count: 0,
498 distinct_values: if distinct { Some(HashSet::new()) } else { None },
499 }
500 }
501}
502
503impl Default for AvgAccumulator {
504 fn default() -> Self {
505 Self::new()
506 }
507}
508
509impl Accumulator for AvgAccumulator {
510 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
511 let Some(value) = value else {
512 return Ok(());
513 };
514 if value.is_null() {
515 return Ok(());
516 }
517 if !distinct_allows(&mut self.distinct_values, &value)? {
518 return Ok(());
519 }
520 let numeric = numeric_to_f64(&value)?;
521 self.sum = Some(self.sum.unwrap_or(0.0) + numeric);
522 self.count += 1;
523 Ok(())
524 }
525
526 fn finalize(&self) -> Result<SqlValue> {
527 if self.count == 0 {
528 return Ok(SqlValue::Null);
529 }
530 let sum = self.sum.unwrap_or(0.0);
531 Ok(SqlValue::Double(sum / self.count as f64))
532 }
533
534 fn state(&self) -> Result<Vec<SqlValue>> {
535 Ok(vec![
536 SqlValue::Double(self.sum.unwrap_or(0.0)),
537 SqlValue::BigInt(self.count as i64),
538 ])
539 }
540
541 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
542 expect_state_arity("avg", state, 2)?;
543 let sum = state_double("avg", &state[0], 0)?;
544 let count = state_bigint("avg", &state[1], 1)?;
545 if count < 0 {
546 return Err(invalid_aggregate_state(
547 "avg",
548 "state count must be non-negative",
549 ));
550 }
551 self.sum = Some(self.sum.unwrap_or(0.0) + sum);
552 self.count = self.count.saturating_add(count as usize);
553 Ok(())
554 }
555
556 fn clone_box(&self) -> Box<dyn Accumulator> {
557 Box::new(self.clone())
558 }
559
560 fn retained_bytes(&self) -> u64 {
561 estimated_distinct_retained_bytes(&self.distinct_values)
562 }
563}
564
565fn numeric_to_f64(value: &SqlValue) -> Result<f64> {
566 match value {
567 SqlValue::Integer(v) => Ok(*v as f64),
568 SqlValue::BigInt(v) => Ok(*v as f64),
569 SqlValue::Float(v) => Ok(*v as f64),
570 SqlValue::Double(v) => Ok(*v),
571 _ => Err(ExecutorError::Evaluation(
572 crate::executor::EvaluationError::TypeMismatch {
573 expected: "numeric".into(),
574 actual: value.type_name().into(),
575 },
576 )),
577 }
578}
579
580fn sum_type_mismatch<T>(expected: &str, actual: &SqlValue) -> Result<T> {
581 Err(ExecutorError::Evaluation(EvaluationError::TypeMismatch {
582 expected: expected.into(),
583 actual: actual.type_name().into(),
584 }))
585}
586
587#[derive(Debug, Clone)]
589pub struct MinMaxAccumulator {
590 value: Option<SqlValue>,
591 is_min: bool,
592 distinct_values: Option<HashSet<Vec<u8>>>,
593}
594
595impl MinMaxAccumulator {
596 pub fn new(is_min: bool) -> Self {
598 Self::with_distinct(is_min, false)
599 }
600
601 pub fn with_distinct(is_min: bool, distinct: bool) -> Self {
602 Self {
603 value: None,
604 is_min,
605 distinct_values: if distinct { Some(HashSet::new()) } else { None },
606 }
607 }
608}
609
610impl Accumulator for MinMaxAccumulator {
611 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
612 let Some(value) = value else {
613 return Ok(());
614 };
615 if value.is_null() {
616 return Ok(());
617 }
618 if !distinct_allows(&mut self.distinct_values, &value)? {
619 return Ok(());
620 }
621
622 match &self.value {
623 None => {
624 self.value = Some(value);
625 }
626 Some(current) => {
627 if std::mem::discriminant(current) != std::mem::discriminant(&value) {
628 return Err(ExecutorError::Evaluation(
629 crate::executor::EvaluationError::TypeMismatch {
630 expected: current.type_name().into(),
631 actual: value.type_name().into(),
632 },
633 ));
634 }
635 let ordering = value.partial_cmp(current).ok_or_else(|| {
636 ExecutorError::Evaluation(crate::executor::EvaluationError::TypeMismatch {
637 expected: current.type_name().into(),
638 actual: value.type_name().into(),
639 })
640 })?;
641 let should_replace = matches!(
642 (self.is_min, ordering),
643 (true, Ordering::Less) | (false, Ordering::Greater)
644 );
645 if should_replace {
646 self.value = Some(value);
647 }
648 }
649 }
650 Ok(())
651 }
652
653 fn finalize(&self) -> Result<SqlValue> {
654 Ok(self.value.clone().unwrap_or(SqlValue::Null))
655 }
656
657 fn state(&self) -> Result<Vec<SqlValue>> {
658 Ok(vec![self.value.clone().unwrap_or(SqlValue::Null)])
659 }
660
661 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
662 expect_state_arity(if self.is_min { "min" } else { "max" }, state, 1)?;
663 if state[0].is_null() {
664 return Ok(());
665 }
666 self.update(Some(state[0].clone()))
667 }
668
669 fn clone_box(&self) -> Box<dyn Accumulator> {
670 Box::new(self.clone())
671 }
672
673 fn retained_bytes(&self) -> u64 {
674 estimated_distinct_retained_bytes(&self.distinct_values).saturating_add(
675 self.value
676 .as_ref()
677 .map(ByteSized::estimated_bytes)
678 .unwrap_or(0),
679 )
680 }
681}
682
683fn compare_ordered_keys(left: &[SqlValue], right: &[SqlValue], specs: &[(bool, bool)]) -> Ordering {
686 for (idx, (asc, nulls_first)) in specs.iter().enumerate() {
687 let left_value = left.get(idx).unwrap_or(&SqlValue::Null);
688 let right_value = right.get(idx).unwrap_or(&SqlValue::Null);
689 let cmp = super::iterator::compare_single(left_value, right_value, *asc, *nulls_first);
690 if cmp != Ordering::Equal {
691 return cmp;
692 }
693 }
694 Ordering::Equal
695}
696
697fn estimated_ordered_string_bytes(values: &[(Vec<SqlValue>, String)]) -> u64 {
698 values.iter().fold(0u64, |total, (keys, value)| {
699 let key_bytes: u64 = keys.iter().map(ByteSized::estimated_bytes).sum();
700 total
701 .saturating_add(key_bytes)
702 .saturating_add(u64::try_from(value.capacity()).unwrap_or(u64::MAX))
703 })
704}
705
706#[derive(Debug, Clone)]
708pub struct GroupConcatAccumulator {
709 values: Vec<String>,
710 separator: String,
711 distinct_values: Option<HashSet<Vec<u8>>>,
712 sort_specs: Vec<(bool, bool)>,
715 ordered_values: Vec<(Vec<SqlValue>, String)>,
716}
717
718impl GroupConcatAccumulator {
719 pub fn new(separator: String) -> Self {
721 Self::with_distinct(separator, false)
722 }
723
724 pub fn with_distinct(separator: String, distinct: bool) -> Self {
725 Self::with_order(separator, distinct, Vec::new())
726 }
727
728 pub fn with_order(separator: String, distinct: bool, sort_specs: Vec<(bool, bool)>) -> Self {
729 Self {
730 values: Vec::new(),
731 separator,
732 distinct_values: if distinct { Some(HashSet::new()) } else { None },
733 sort_specs,
734 ordered_values: Vec::new(),
735 }
736 }
737}
738
739impl Accumulator for GroupConcatAccumulator {
740 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
741 if !self.sort_specs.is_empty() {
742 return Err(invalid_aggregate_state(
743 "group_concat",
744 "ordered aggregation requires the sort keys of every row",
745 ));
746 }
747 let Some(value) = value else {
748 return Ok(());
749 };
750 match value {
751 SqlValue::Null => Ok(()),
752 SqlValue::Text(text) => {
753 let value = SqlValue::Text(text.clone());
754 if !distinct_allows(&mut self.distinct_values, &value)? {
755 return Ok(());
756 }
757 self.values.push(text);
758 Ok(())
759 }
760 other => Err(ExecutorError::Evaluation(
761 crate::executor::EvaluationError::TypeMismatch {
762 expected: "Text".into(),
763 actual: other.type_name().into(),
764 },
765 )),
766 }
767 }
768
769 fn update_ordered(&mut self, value: Option<SqlValue>, sort_keys: &[SqlValue]) -> Result<()> {
770 if self.sort_specs.is_empty() {
771 return self.update(value);
772 }
773 let Some(value) = value else {
774 return Ok(());
775 };
776 match value {
777 SqlValue::Null => Ok(()),
778 SqlValue::Text(text) => {
779 let value = SqlValue::Text(text.clone());
780 if !distinct_allows(&mut self.distinct_values, &value)? {
781 return Ok(());
782 }
783 self.ordered_values.push((sort_keys.to_vec(), text));
784 Ok(())
785 }
786 other => Err(ExecutorError::Evaluation(
787 crate::executor::EvaluationError::TypeMismatch {
788 expected: "Text".into(),
789 actual: other.type_name().into(),
790 },
791 )),
792 }
793 }
794
795 fn finalize(&self) -> Result<SqlValue> {
796 if !self.sort_specs.is_empty() {
797 if self.ordered_values.is_empty() {
798 return Ok(SqlValue::Null);
799 }
800 let mut sorted = self.ordered_values.clone();
801 sorted.sort_by(|left, right| compare_ordered_keys(&left.0, &right.0, &self.sort_specs));
802 let joined = sorted
803 .into_iter()
804 .map(|(_, value)| value)
805 .collect::<Vec<_>>()
806 .join(&self.separator);
807 return Ok(SqlValue::Text(joined));
808 }
809 if self.values.is_empty() {
810 return Ok(SqlValue::Null);
811 }
812 Ok(SqlValue::Text(self.values.join(&self.separator)))
813 }
814
815 fn state(&self) -> Result<Vec<SqlValue>> {
816 if !self.sort_specs.is_empty() {
817 return Err(invalid_aggregate_state(
818 "group_concat",
819 "ordered aggregation cannot produce partial state",
820 ));
821 }
822 Ok(vec![
823 if self.values.is_empty() {
824 SqlValue::Null
825 } else {
826 SqlValue::Text(self.values.join(&self.separator))
827 },
828 SqlValue::Text(self.separator.clone()),
829 ])
830 }
831
832 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
833 if !self.sort_specs.is_empty() {
834 return Err(invalid_aggregate_state(
835 "group_concat",
836 "ordered aggregation cannot merge partial state",
837 ));
838 }
839 expect_state_arity("group_concat", state, 2)?;
840 let separator = state_text("group_concat", &state[1], 1)?;
841 if separator != self.separator {
842 return Err(invalid_aggregate_state(
843 "group_concat",
844 "state separator differs from accumulator separator",
845 ));
846 }
847 match &state[0] {
848 SqlValue::Null => Ok(()),
849 SqlValue::Text(text) => {
850 self.values.push(text.clone());
851 Ok(())
852 }
853 other => Err(invalid_aggregate_state(
854 "group_concat",
855 format!(
856 "state value 0 expected Text or Null, got {}",
857 other.type_name()
858 ),
859 )),
860 }
861 }
862
863 fn clone_box(&self) -> Box<dyn Accumulator> {
864 Box::new(self.clone())
865 }
866
867 fn retained_bytes(&self) -> u64 {
868 estimated_distinct_retained_bytes(&self.distinct_values)
869 .saturating_add(estimated_string_collection_bytes(
870 &self.values,
871 self.values.capacity(),
872 self.separator.capacity(),
873 ))
874 .saturating_add(estimated_ordered_string_bytes(&self.ordered_values))
875 }
876}
877
878#[derive(Debug, Clone)]
880pub struct StringAggAccumulator {
881 values: Vec<String>,
882 separator: String,
883 distinct_values: Option<HashSet<Vec<u8>>>,
884 sort_specs: Vec<(bool, bool)>,
886 ordered_values: Vec<(Vec<SqlValue>, String)>,
887}
888
889impl StringAggAccumulator {
890 pub fn new(separator: String) -> Self {
892 Self::with_distinct(separator, false)
893 }
894
895 pub fn with_distinct(separator: String, distinct: bool) -> Self {
896 Self::with_order(separator, distinct, Vec::new())
897 }
898
899 pub fn with_order(separator: String, distinct: bool, sort_specs: Vec<(bool, bool)>) -> Self {
900 Self {
901 values: Vec::new(),
902 separator,
903 distinct_values: if distinct { Some(HashSet::new()) } else { None },
904 sort_specs,
905 ordered_values: Vec::new(),
906 }
907 }
908}
909
910impl Accumulator for StringAggAccumulator {
911 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
912 if !self.sort_specs.is_empty() {
913 return Err(invalid_aggregate_state(
914 "string_agg",
915 "ordered aggregation requires the sort keys of every row",
916 ));
917 }
918 let Some(value) = value else {
919 return Ok(());
920 };
921 match value {
922 SqlValue::Null => Ok(()),
923 SqlValue::Text(s) => {
924 let value = SqlValue::Text(s.clone());
925 if !distinct_allows(&mut self.distinct_values, &value)? {
926 return Ok(());
927 }
928 self.values.push(s);
929 Ok(())
930 }
931 other => Err(ExecutorError::Evaluation(
932 crate::executor::EvaluationError::TypeMismatch {
933 expected: "Text".into(),
934 actual: other.type_name().into(),
935 },
936 )),
937 }
938 }
939
940 fn update_ordered(&mut self, value: Option<SqlValue>, sort_keys: &[SqlValue]) -> Result<()> {
941 if self.sort_specs.is_empty() {
942 return self.update(value);
943 }
944 let Some(value) = value else {
945 return Ok(());
946 };
947 match value {
948 SqlValue::Null => Ok(()),
949 SqlValue::Text(s) => {
950 let value = SqlValue::Text(s.clone());
951 if !distinct_allows(&mut self.distinct_values, &value)? {
952 return Ok(());
953 }
954 self.ordered_values.push((sort_keys.to_vec(), s));
955 Ok(())
956 }
957 other => Err(ExecutorError::Evaluation(
958 crate::executor::EvaluationError::TypeMismatch {
959 expected: "Text".into(),
960 actual: other.type_name().into(),
961 },
962 )),
963 }
964 }
965
966 fn finalize(&self) -> Result<SqlValue> {
967 if !self.sort_specs.is_empty() {
968 if self.ordered_values.is_empty() {
969 return Ok(SqlValue::Null);
970 }
971 let mut sorted = self.ordered_values.clone();
972 sorted.sort_by(|left, right| compare_ordered_keys(&left.0, &right.0, &self.sort_specs));
973 let joined = sorted
974 .into_iter()
975 .map(|(_, value)| value)
976 .collect::<Vec<_>>()
977 .join(&self.separator);
978 return Ok(SqlValue::Text(joined));
979 }
980 if self.values.is_empty() {
981 return Ok(SqlValue::Null);
982 }
983 Ok(SqlValue::Text(self.values.join(&self.separator)))
984 }
985
986 fn state(&self) -> Result<Vec<SqlValue>> {
987 if !self.sort_specs.is_empty() {
988 return Err(invalid_aggregate_state(
989 "string_agg",
990 "ordered aggregation cannot produce partial state",
991 ));
992 }
993 Ok(vec![
994 if self.values.is_empty() {
995 SqlValue::Null
996 } else {
997 SqlValue::Text(self.values.join(&self.separator))
998 },
999 SqlValue::Text(self.separator.clone()),
1000 ])
1001 }
1002
1003 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
1004 if !self.sort_specs.is_empty() {
1005 return Err(invalid_aggregate_state(
1006 "string_agg",
1007 "ordered aggregation cannot merge partial state",
1008 ));
1009 }
1010 expect_state_arity("string_agg", state, 2)?;
1011 let separator = state_text("string_agg", &state[1], 1)?;
1012 if separator != self.separator {
1013 return Err(invalid_aggregate_state(
1014 "string_agg",
1015 "state separator differs from accumulator separator",
1016 ));
1017 }
1018 match &state[0] {
1019 SqlValue::Null => Ok(()),
1020 SqlValue::Text(text) => {
1021 self.values.push(text.clone());
1022 Ok(())
1023 }
1024 other => Err(invalid_aggregate_state(
1025 "string_agg",
1026 format!(
1027 "state value 0 expected Text or Null, got {}",
1028 other.type_name()
1029 ),
1030 )),
1031 }
1032 }
1033
1034 fn clone_box(&self) -> Box<dyn Accumulator> {
1035 Box::new(self.clone())
1036 }
1037
1038 fn retained_bytes(&self) -> u64 {
1039 estimated_distinct_retained_bytes(&self.distinct_values)
1040 .saturating_add(estimated_string_collection_bytes(
1041 &self.values,
1042 self.values.capacity(),
1043 self.separator.capacity(),
1044 ))
1045 .saturating_add(estimated_ordered_string_bytes(&self.ordered_values))
1046 }
1047}
1048
1049#[derive(Debug, Clone)]
1057pub struct PercentileDiscAccumulator {
1058 fraction: f64,
1059 sort_specs: Vec<(bool, bool)>,
1060 values: Vec<(Vec<SqlValue>, SqlValue)>,
1061}
1062
1063impl PercentileDiscAccumulator {
1064 pub fn new(fraction: f64, sort_specs: Vec<(bool, bool)>) -> Self {
1065 let sort_specs = if sort_specs.is_empty() {
1066 vec![(true, false)]
1067 } else {
1068 sort_specs
1069 };
1070 Self {
1071 fraction,
1072 sort_specs,
1073 values: Vec::new(),
1074 }
1075 }
1076}
1077
1078impl Accumulator for PercentileDiscAccumulator {
1079 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1080 let keys = match &value {
1083 Some(inner) if !inner.is_null() => vec![inner.clone()],
1084 _ => return Ok(()),
1085 };
1086 self.update_ordered(value, &keys)
1087 }
1088
1089 fn update_ordered(&mut self, value: Option<SqlValue>, sort_keys: &[SqlValue]) -> Result<()> {
1090 let Some(value) = value else {
1091 return Ok(());
1092 };
1093 if value.is_null() {
1094 return Ok(());
1095 }
1096 self.values.push((sort_keys.to_vec(), value));
1097 Ok(())
1098 }
1099
1100 fn finalize(&self) -> Result<SqlValue> {
1101 if self.values.is_empty() {
1102 return Ok(SqlValue::Null);
1103 }
1104 let mut sorted = self.values.clone();
1105 sorted.sort_by(|left, right| compare_ordered_keys(&left.0, &right.0, &self.sort_specs));
1106 let count = sorted.len();
1107 let index = (self.fraction * count as f64).ceil() as usize;
1108 let index = index.saturating_sub(1).min(count - 1);
1109 Ok(sorted[index].1.clone())
1110 }
1111
1112 fn state(&self) -> Result<Vec<SqlValue>> {
1113 Err(invalid_aggregate_state(
1114 "percentile_disc",
1115 "ordered-set aggregation cannot produce partial state",
1116 ))
1117 }
1118
1119 fn merge(&mut self, _state: &[SqlValue]) -> Result<()> {
1120 Err(invalid_aggregate_state(
1121 "percentile_disc",
1122 "ordered-set aggregation cannot merge partial state",
1123 ))
1124 }
1125
1126 fn clone_box(&self) -> Box<dyn Accumulator> {
1127 Box::new(self.clone())
1128 }
1129
1130 fn retained_bytes(&self) -> u64 {
1131 self.values.iter().fold(0u64, |total, (keys, value)| {
1132 let key_bytes: u64 = keys.iter().map(ByteSized::estimated_bytes).sum();
1133 total
1134 .saturating_add(key_bytes)
1135 .saturating_add(value.estimated_bytes())
1136 })
1137 }
1138}
1139
1140#[derive(Debug, Clone, Copy)]
1141enum StatisticsKind {
1142 Variance(bool),
1143 Stddev(bool),
1144 Covariance(bool),
1145 Corr,
1146 RegrCount,
1147 RegrAvgX,
1148 RegrAvgY,
1149 RegrSxx,
1150 RegrSyy,
1151 RegrSxy,
1152 RegrSlope,
1153 RegrIntercept,
1154 RegrR2,
1155}
1156
1157#[derive(Debug, Clone)]
1159struct StatisticsAccumulator {
1160 kind: StatisticsKind,
1161 count: u64,
1162 mean_x: f64,
1163 mean_y: f64,
1164 m2_x: f64,
1165 m2_y: f64,
1166 co_moment: f64,
1167}
1168
1169impl StatisticsAccumulator {
1170 fn new(kind: StatisticsKind) -> Self {
1171 Self {
1172 kind,
1173 count: 0,
1174 mean_x: 0.0,
1175 mean_y: 0.0,
1176 m2_x: 0.0,
1177 m2_y: 0.0,
1178 co_moment: 0.0,
1179 }
1180 }
1181
1182 fn update_pair(&mut self, y: f64, x: f64) {
1183 self.count += 1;
1184 let count = self.count as f64;
1185 let dx = x - self.mean_x;
1186 let dy = y - self.mean_y;
1187 self.mean_x += dx / count;
1188 self.mean_y += dy / count;
1189 self.m2_x += dx * (x - self.mean_x);
1190 self.m2_y += dy * (y - self.mean_y);
1191 self.co_moment += dx * (y - self.mean_y);
1192 }
1193
1194 fn null_or(&self, value: f64) -> SqlValue {
1195 if self.count == 0 {
1196 SqlValue::Null
1197 } else {
1198 SqlValue::Double(value)
1199 }
1200 }
1201}
1202
1203impl Accumulator for StatisticsAccumulator {
1204 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1205 let Some(value) = value.filter(|value| !value.is_null()) else {
1206 return Ok(());
1207 };
1208 let value = numeric_to_f64(&value)?;
1209 self.update_pair(value, value);
1210 Ok(())
1211 }
1212
1213 fn update_values(&mut self, values: &[SqlValue]) -> Result<()> {
1214 if values.len() == 1 {
1215 return self.update(values.first().cloned());
1216 }
1217 let [y, x] = values else {
1218 return Err(invalid_aggregate_state(
1219 "statistics",
1220 format!("expected one or two inputs, got {}", values.len()),
1221 ));
1222 };
1223 if y.is_null() || x.is_null() {
1224 return Ok(());
1225 }
1226 self.update_pair(numeric_to_f64(y)?, numeric_to_f64(x)?);
1227 Ok(())
1228 }
1229
1230 fn state(&self) -> Result<Vec<SqlValue>> {
1231 Ok(vec![
1232 SqlValue::BigInt(self.count as i64),
1233 SqlValue::Double(self.mean_x),
1234 SqlValue::Double(self.mean_y),
1235 SqlValue::Double(self.m2_x),
1236 SqlValue::Double(self.m2_y),
1237 SqlValue::Double(self.co_moment),
1238 ])
1239 }
1240
1241 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
1242 expect_state_arity("statistics", state, 6)?;
1243 let other_count = state_bigint("statistics", &state[0], 0)?;
1244 if other_count < 0 {
1245 return Err(invalid_aggregate_state(
1246 "statistics",
1247 "state count must be non-negative",
1248 ));
1249 }
1250 let other_count = other_count as u64;
1251 if other_count == 0 {
1252 return Ok(());
1253 }
1254 let other_mean_x = state_double("statistics", &state[1], 1)?;
1255 let other_mean_y = state_double("statistics", &state[2], 2)?;
1256 let other_m2_x = state_double("statistics", &state[3], 3)?;
1257 let other_m2_y = state_double("statistics", &state[4], 4)?;
1258 let other_co_moment = state_double("statistics", &state[5], 5)?;
1259 if self.count == 0 {
1260 self.count = other_count;
1261 self.mean_x = other_mean_x;
1262 self.mean_y = other_mean_y;
1263 self.m2_x = other_m2_x;
1264 self.m2_y = other_m2_y;
1265 self.co_moment = other_co_moment;
1266 return Ok(());
1267 }
1268 let total = self.count + other_count;
1269 let left = self.count as f64;
1270 let right = other_count as f64;
1271 let total_f = total as f64;
1272 let dx = other_mean_x - self.mean_x;
1273 let dy = other_mean_y - self.mean_y;
1274 self.m2_x += other_m2_x + dx * dx * left * right / total_f;
1275 self.m2_y += other_m2_y + dy * dy * left * right / total_f;
1276 self.co_moment += other_co_moment + dx * dy * left * right / total_f;
1277 self.mean_x += dx * right / total_f;
1278 self.mean_y += dy * right / total_f;
1279 self.count = total;
1280 Ok(())
1281 }
1282
1283 fn finalize(&self) -> Result<SqlValue> {
1284 let sample_divisor = || (self.count > 1).then(|| (self.count - 1) as f64);
1285 let population_divisor = || (self.count > 0).then_some(self.count as f64);
1286 let value = match self.kind {
1287 StatisticsKind::Variance(sample) | StatisticsKind::Stddev(sample) => {
1288 let divisor = if sample {
1289 sample_divisor()
1290 } else {
1291 population_divisor()
1292 };
1293 let Some(divisor) = divisor else {
1294 return Ok(SqlValue::Null);
1295 };
1296 let variance = self.m2_x / divisor;
1297 if matches!(self.kind, StatisticsKind::Stddev(_)) {
1298 variance.sqrt()
1299 } else {
1300 variance
1301 }
1302 }
1303 StatisticsKind::Covariance(sample) => {
1304 let divisor = if sample {
1305 sample_divisor()
1306 } else {
1307 population_divisor()
1308 };
1309 let Some(divisor) = divisor else {
1310 return Ok(SqlValue::Null);
1311 };
1312 self.co_moment / divisor
1313 }
1314 StatisticsKind::Corr => {
1315 let divisor = (self.m2_x * self.m2_y).sqrt();
1316 if self.count == 0 || divisor == 0.0 {
1317 return Ok(SqlValue::Null);
1318 }
1319 self.co_moment / divisor
1320 }
1321 StatisticsKind::RegrCount => return Ok(SqlValue::BigInt(self.count as i64)),
1322 StatisticsKind::RegrAvgX => return Ok(self.null_or(self.mean_x)),
1323 StatisticsKind::RegrAvgY => return Ok(self.null_or(self.mean_y)),
1324 StatisticsKind::RegrSxx => return Ok(self.null_or(self.m2_x)),
1325 StatisticsKind::RegrSyy => return Ok(self.null_or(self.m2_y)),
1326 StatisticsKind::RegrSxy => return Ok(self.null_or(self.co_moment)),
1327 StatisticsKind::RegrSlope => {
1328 if self.count == 0 || self.m2_x == 0.0 {
1329 return Ok(SqlValue::Null);
1330 }
1331 self.co_moment / self.m2_x
1332 }
1333 StatisticsKind::RegrIntercept => {
1334 if self.count == 0 || self.m2_x == 0.0 {
1335 return Ok(SqlValue::Null);
1336 }
1337 self.mean_y - (self.co_moment / self.m2_x) * self.mean_x
1338 }
1339 StatisticsKind::RegrR2 => {
1340 if self.count == 0 || self.m2_x == 0.0 {
1341 return Ok(SqlValue::Null);
1342 }
1343 if self.m2_y == 0.0 {
1344 1.0
1345 } else {
1346 self.co_moment * self.co_moment / (self.m2_x * self.m2_y)
1347 }
1348 }
1349 };
1350 Ok(SqlValue::Double(value))
1351 }
1352
1353 fn clone_box(&self) -> Box<dyn Accumulator> {
1354 Box::new(self.clone())
1355 }
1356}
1357
1358#[derive(Debug, Clone)]
1359struct PercentileContAccumulator {
1360 fraction: f64,
1361 ascending: bool,
1362 values: Vec<f64>,
1363}
1364
1365impl PercentileContAccumulator {
1366 fn new(fraction: f64, ascending: bool) -> Self {
1367 Self {
1368 fraction,
1369 ascending,
1370 values: Vec::new(),
1371 }
1372 }
1373}
1374
1375impl Accumulator for PercentileContAccumulator {
1376 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1377 let Some(value) = value.filter(|value| !value.is_null()) else {
1378 return Ok(());
1379 };
1380 self.values.push(numeric_to_f64(&value)?);
1381 Ok(())
1382 }
1383
1384 fn state(&self) -> Result<Vec<SqlValue>> {
1385 Err(invalid_aggregate_state(
1386 "percentile_cont",
1387 "ordered aggregation cannot produce partial state",
1388 ))
1389 }
1390
1391 fn merge(&mut self, _state: &[SqlValue]) -> Result<()> {
1392 Err(invalid_aggregate_state(
1393 "percentile_cont",
1394 "ordered aggregation cannot merge partial state",
1395 ))
1396 }
1397
1398 fn finalize(&self) -> Result<SqlValue> {
1399 if self.values.is_empty() {
1400 return Ok(SqlValue::Null);
1401 }
1402 let mut values = self.values.clone();
1403 values.sort_by(f64::total_cmp);
1404 if !self.ascending {
1405 values.reverse();
1406 }
1407 let position = self.fraction * (values.len() - 1) as f64;
1408 let lower = position.floor() as usize;
1409 let upper = position.ceil() as usize;
1410 if lower == upper {
1411 return Ok(SqlValue::Double(values[lower]));
1412 }
1413 let weight = position - lower as f64;
1414 Ok(SqlValue::Double(
1415 values[lower] + (values[upper] - values[lower]) * weight,
1416 ))
1417 }
1418
1419 fn clone_box(&self) -> Box<dyn Accumulator> {
1420 Box::new(self.clone())
1421 }
1422
1423 fn retained_bytes(&self) -> u64 {
1424 u64::try_from(
1425 self.values
1426 .capacity()
1427 .saturating_mul(std::mem::size_of::<f64>()),
1428 )
1429 .unwrap_or(u64::MAX)
1430 }
1431}
1432
1433#[derive(Debug, Clone)]
1434struct ModeAccumulator {
1435 sort_specs: Vec<(bool, bool)>,
1436 values: Vec<(Vec<SqlValue>, SqlValue)>,
1437}
1438
1439impl ModeAccumulator {
1440 fn new(sort_specs: Vec<(bool, bool)>) -> Self {
1441 Self {
1442 sort_specs: if sort_specs.is_empty() {
1443 vec![(true, false)]
1444 } else {
1445 sort_specs
1446 },
1447 values: Vec::new(),
1448 }
1449 }
1450}
1451
1452impl Accumulator for ModeAccumulator {
1453 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1454 let Some(value) = value.filter(|value| !value.is_null()) else {
1455 return Ok(());
1456 };
1457 self.values.push((vec![value.clone()], value));
1458 Ok(())
1459 }
1460
1461 fn update_ordered(&mut self, value: Option<SqlValue>, sort_keys: &[SqlValue]) -> Result<()> {
1462 let Some(value) = value.filter(|value| !value.is_null()) else {
1463 return Ok(());
1464 };
1465 self.values.push((sort_keys.to_vec(), value));
1466 Ok(())
1467 }
1468
1469 fn state(&self) -> Result<Vec<SqlValue>> {
1470 Err(invalid_aggregate_state(
1471 "mode",
1472 "ordered aggregation cannot produce partial state",
1473 ))
1474 }
1475
1476 fn merge(&mut self, _state: &[SqlValue]) -> Result<()> {
1477 Err(invalid_aggregate_state(
1478 "mode",
1479 "ordered aggregation cannot merge partial state",
1480 ))
1481 }
1482
1483 fn finalize(&self) -> Result<SqlValue> {
1484 if self.values.is_empty() {
1485 return Ok(SqlValue::Null);
1486 }
1487 let mut values = self.values.clone();
1488 values.sort_by(|left, right| compare_ordered_keys(&left.0, &right.0, &self.sort_specs));
1489 let mut best_start = 0;
1490 let mut best_count = 0;
1491 let mut start = 0;
1492 while start < values.len() {
1493 let mut end = start + 1;
1494 while end < values.len()
1495 && compare_ordered_keys(&values[start].0, &values[end].0, &self.sort_specs)
1496 == Ordering::Equal
1497 {
1498 end += 1;
1499 }
1500 if end - start > best_count {
1501 best_start = start;
1502 best_count = end - start;
1503 }
1504 start = end;
1505 }
1506 Ok(values[best_start].1.clone())
1507 }
1508
1509 fn clone_box(&self) -> Box<dyn Accumulator> {
1510 Box::new(self.clone())
1511 }
1512
1513 fn retained_bytes(&self) -> u64 {
1514 self.values.iter().fold(0, |total, (keys, value)| {
1515 total
1516 .saturating_add(keys.iter().map(ByteSized::estimated_bytes).sum::<u64>())
1517 .saturating_add(value.estimated_bytes())
1518 })
1519 }
1520}
1521
1522#[derive(Debug, Clone, Copy)]
1523enum ValueKind {
1524 Any,
1525 First,
1526 Last,
1527 ArgMin,
1528 ArgMax,
1529}
1530
1531#[derive(Debug, Clone)]
1532struct ValueAccumulator {
1533 kind: ValueKind,
1534 sort_specs: Vec<(bool, bool)>,
1535 chosen: Option<(Vec<SqlValue>, SqlValue)>,
1536}
1537
1538impl ValueAccumulator {
1539 fn new(kind: ValueKind, sort_specs: Vec<(bool, bool)>) -> Self {
1540 Self {
1541 kind,
1542 sort_specs,
1543 chosen: None,
1544 }
1545 }
1546
1547 fn select(&mut self, keys: Vec<SqlValue>, value: SqlValue) {
1548 let replace = match &self.chosen {
1549 None => true,
1550 Some((current, _)) => match self.kind {
1551 ValueKind::Any | ValueKind::First => false,
1552 ValueKind::Last => true,
1553 ValueKind::ArgMin => {
1554 compare_ordered_keys(&keys, current, &[(true, false)]) == Ordering::Less
1555 }
1556 ValueKind::ArgMax => {
1557 compare_ordered_keys(&keys, current, &[(true, false)]) == Ordering::Greater
1558 }
1559 },
1560 };
1561 if replace {
1562 self.chosen = Some((keys, value));
1563 }
1564 }
1565}
1566
1567impl Accumulator for ValueAccumulator {
1568 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1569 let Some(value) = value else {
1570 return Ok(());
1571 };
1572 if matches!(self.kind, ValueKind::Any) && value.is_null() {
1573 return Ok(());
1574 }
1575 self.select(Vec::new(), value);
1576 Ok(())
1577 }
1578
1579 fn update_values(&mut self, values: &[SqlValue]) -> Result<()> {
1580 if matches!(self.kind, ValueKind::ArgMin | ValueKind::ArgMax) {
1581 let [value, key] = values else {
1582 return Err(invalid_aggregate_state(
1583 "arg_min/max",
1584 format!("expected two inputs, got {}", values.len()),
1585 ));
1586 };
1587 if value.is_null() || key.is_null() {
1588 return Ok(());
1589 }
1590 self.select(vec![key.clone()], value.clone());
1591 return Ok(());
1592 }
1593 self.update(values.first().cloned())
1594 }
1595
1596 fn update_ordered(&mut self, value: Option<SqlValue>, sort_keys: &[SqlValue]) -> Result<()> {
1597 let Some(value) = value else {
1598 return Ok(());
1599 };
1600 match &self.chosen {
1601 None => self.chosen = Some((sort_keys.to_vec(), value)),
1602 Some((keys, _)) => {
1603 let ordering = compare_ordered_keys(sort_keys, keys, &self.sort_specs);
1604 let replace = matches!(
1605 (self.kind, ordering),
1606 (ValueKind::First, Ordering::Less) | (ValueKind::Last, Ordering::Greater)
1607 );
1608 if replace {
1609 self.chosen = Some((sort_keys.to_vec(), value));
1610 }
1611 }
1612 }
1613 Ok(())
1614 }
1615
1616 fn state(&self) -> Result<Vec<SqlValue>> {
1617 Ok(vec![
1618 self.chosen
1619 .as_ref()
1620 .map(|(_, value)| value.clone())
1621 .unwrap_or(SqlValue::Null),
1622 ])
1623 }
1624
1625 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
1626 expect_state_arity("value aggregate", state, 1)?;
1627 self.update(Some(state[0].clone()))
1628 }
1629
1630 fn finalize(&self) -> Result<SqlValue> {
1631 Ok(self
1632 .chosen
1633 .as_ref()
1634 .map(|(_, value)| value.clone())
1635 .unwrap_or(SqlValue::Null))
1636 }
1637
1638 fn clone_box(&self) -> Box<dyn Accumulator> {
1639 Box::new(self.clone())
1640 }
1641
1642 fn retained_bytes(&self) -> u64 {
1643 self.chosen
1644 .as_ref()
1645 .map(|(keys, value)| {
1646 keys.iter()
1647 .map(ByteSized::estimated_bytes)
1648 .sum::<u64>()
1649 .saturating_add(value.estimated_bytes())
1650 })
1651 .unwrap_or(0)
1652 }
1653}
1654
1655#[derive(Debug, Clone, Copy)]
1656enum BitKind {
1657 And,
1658 Or,
1659 Xor,
1660}
1661
1662#[derive(Debug, Clone)]
1663struct BitAccumulator {
1664 kind: BitKind,
1665 value: Option<SqlValue>,
1666}
1667
1668impl BitAccumulator {
1669 fn new(kind: BitKind) -> Self {
1670 Self { kind, value: None }
1671 }
1672}
1673
1674impl Accumulator for BitAccumulator {
1675 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1676 let Some(value) = value.filter(|value| !value.is_null()) else {
1677 return Ok(());
1678 };
1679 self.value = Some(match (self.value.take(), value) {
1680 (None, value) => value,
1681 (Some(SqlValue::Integer(left)), SqlValue::Integer(right)) => {
1682 SqlValue::Integer(match self.kind {
1683 BitKind::And => left & right,
1684 BitKind::Or => left | right,
1685 BitKind::Xor => left ^ right,
1686 })
1687 }
1688 (Some(SqlValue::BigInt(left)), SqlValue::BigInt(right)) => {
1689 SqlValue::BigInt(match self.kind {
1690 BitKind::And => left & right,
1691 BitKind::Or => left | right,
1692 BitKind::Xor => left ^ right,
1693 })
1694 }
1695 (Some(left), right) => {
1696 return Err(ExecutorError::Evaluation(EvaluationError::TypeMismatch {
1697 expected: "matching integer types".into(),
1698 actual: format!("{} vs {}", left.type_name(), right.type_name()),
1699 }));
1700 }
1701 });
1702 Ok(())
1703 }
1704
1705 fn state(&self) -> Result<Vec<SqlValue>> {
1706 Ok(vec![self.value.clone().unwrap_or(SqlValue::Null)])
1707 }
1708
1709 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
1710 expect_state_arity("bit aggregate", state, 1)?;
1711 self.update(Some(state[0].clone()))
1712 }
1713
1714 fn finalize(&self) -> Result<SqlValue> {
1715 Ok(self.value.clone().unwrap_or(SqlValue::Null))
1716 }
1717
1718 fn clone_box(&self) -> Box<dyn Accumulator> {
1719 Box::new(self.clone())
1720 }
1721}
1722
1723#[derive(Debug, Clone)]
1724struct BoolAccumulator {
1725 and: bool,
1726 value: Option<bool>,
1727}
1728
1729impl BoolAccumulator {
1730 fn new(and: bool) -> Self {
1731 Self { and, value: None }
1732 }
1733}
1734
1735impl Accumulator for BoolAccumulator {
1736 fn update(&mut self, value: Option<SqlValue>) -> Result<()> {
1737 match value {
1738 None | Some(SqlValue::Null) => Ok(()),
1739 Some(SqlValue::Boolean(value)) => {
1740 self.value = Some(match self.value {
1741 None => value,
1742 Some(current) if self.and => current && value,
1743 Some(current) => current || value,
1744 });
1745 Ok(())
1746 }
1747 Some(other) => Err(ExecutorError::Evaluation(EvaluationError::TypeMismatch {
1748 expected: "Boolean".into(),
1749 actual: other.type_name().into(),
1750 })),
1751 }
1752 }
1753
1754 fn state(&self) -> Result<Vec<SqlValue>> {
1755 Ok(vec![
1756 self.value.map(SqlValue::Boolean).unwrap_or(SqlValue::Null),
1757 ])
1758 }
1759
1760 fn merge(&mut self, state: &[SqlValue]) -> Result<()> {
1761 expect_state_arity("boolean aggregate", state, 1)?;
1762 self.update(Some(state[0].clone()))
1763 }
1764
1765 fn finalize(&self) -> Result<SqlValue> {
1766 Ok(self.value.map(SqlValue::Boolean).unwrap_or(SqlValue::Null))
1767 }
1768
1769 fn clone_box(&self) -> Box<dyn Accumulator> {
1770 Box::new(self.clone())
1771 }
1772}
1773
1774pub fn create_accumulator(function: &AggregateFunction, distinct: bool) -> Box<dyn Accumulator> {
1776 match function {
1777 AggregateFunction::Count => Box::new(CountAccumulator::new(distinct)),
1778 AggregateFunction::Sum => Box::new(SumAccumulator::with_distinct(distinct)),
1779 AggregateFunction::Total => Box::new(TotalAccumulator::new()),
1780 AggregateFunction::Avg => Box::new(AvgAccumulator::with_distinct(distinct)),
1781 AggregateFunction::Min => Box::new(MinMaxAccumulator::with_distinct(true, distinct)),
1782 AggregateFunction::Max => Box::new(MinMaxAccumulator::with_distinct(false, distinct)),
1783 AggregateFunction::GroupConcat { separator } => {
1784 let sep = separator.clone().unwrap_or_else(|| ",".to_string());
1785 Box::new(GroupConcatAccumulator::with_distinct(sep, distinct))
1786 }
1787 AggregateFunction::StringAgg { separator } => {
1788 let sep = separator.clone().unwrap_or_else(|| ",".to_string());
1789 Box::new(StringAggAccumulator::with_distinct(sep, distinct))
1790 }
1791 AggregateFunction::PercentileDisc { fraction } => {
1792 Box::new(PercentileDiscAccumulator::new(*fraction, Vec::new()))
1793 }
1794 AggregateFunction::PercentileCont { fraction }
1795 | AggregateFunction::QuantileCont { fraction } => {
1796 Box::new(PercentileContAccumulator::new(*fraction, true))
1797 }
1798 AggregateFunction::Variance { sample } => Box::new(StatisticsAccumulator::new(
1799 StatisticsKind::Variance(*sample),
1800 )),
1801 AggregateFunction::Stddev { sample } => {
1802 Box::new(StatisticsAccumulator::new(StatisticsKind::Stddev(*sample)))
1803 }
1804 AggregateFunction::Covariance { sample } => Box::new(StatisticsAccumulator::new(
1805 StatisticsKind::Covariance(*sample),
1806 )),
1807 AggregateFunction::Corr => Box::new(StatisticsAccumulator::new(StatisticsKind::Corr)),
1808 AggregateFunction::Median => Box::new(PercentileContAccumulator::new(0.5, true)),
1809 AggregateFunction::Mode => Box::new(ModeAccumulator::new(Vec::new())),
1810 AggregateFunction::RegrCount => {
1811 Box::new(StatisticsAccumulator::new(StatisticsKind::RegrCount))
1812 }
1813 AggregateFunction::RegrAvgX => {
1814 Box::new(StatisticsAccumulator::new(StatisticsKind::RegrAvgX))
1815 }
1816 AggregateFunction::RegrAvgY => {
1817 Box::new(StatisticsAccumulator::new(StatisticsKind::RegrAvgY))
1818 }
1819 AggregateFunction::RegrSxx => Box::new(StatisticsAccumulator::new(StatisticsKind::RegrSxx)),
1820 AggregateFunction::RegrSyy => Box::new(StatisticsAccumulator::new(StatisticsKind::RegrSyy)),
1821 AggregateFunction::RegrSxy => Box::new(StatisticsAccumulator::new(StatisticsKind::RegrSxy)),
1822 AggregateFunction::RegrSlope => {
1823 Box::new(StatisticsAccumulator::new(StatisticsKind::RegrSlope))
1824 }
1825 AggregateFunction::RegrIntercept => {
1826 Box::new(StatisticsAccumulator::new(StatisticsKind::RegrIntercept))
1827 }
1828 AggregateFunction::RegrR2 => Box::new(StatisticsAccumulator::new(StatisticsKind::RegrR2)),
1829 AggregateFunction::AnyValue => Box::new(ValueAccumulator::new(ValueKind::Any, Vec::new())),
1830 AggregateFunction::First => Box::new(ValueAccumulator::new(ValueKind::First, Vec::new())),
1831 AggregateFunction::Last => Box::new(ValueAccumulator::new(ValueKind::Last, Vec::new())),
1832 AggregateFunction::ArgMin => Box::new(ValueAccumulator::new(ValueKind::ArgMin, Vec::new())),
1833 AggregateFunction::ArgMax => Box::new(ValueAccumulator::new(ValueKind::ArgMax, Vec::new())),
1834 AggregateFunction::BitAnd => Box::new(BitAccumulator::new(BitKind::And)),
1835 AggregateFunction::BitOr => Box::new(BitAccumulator::new(BitKind::Or)),
1836 AggregateFunction::BitXor => Box::new(BitAccumulator::new(BitKind::Xor)),
1837 AggregateFunction::BoolAnd => Box::new(BoolAccumulator::new(true)),
1838 AggregateFunction::BoolOr => Box::new(BoolAccumulator::new(false)),
1839 }
1840}
1841
1842fn aggregate_sort_specs(aggregate: &AggregateExpr) -> Vec<(bool, bool)> {
1843 aggregate
1844 .order_by
1845 .iter()
1846 .map(|sort| (sort.asc, sort.nulls_first))
1847 .collect()
1848}
1849
1850pub fn create_accumulator_for_aggregate(aggregate: &AggregateExpr) -> Box<dyn Accumulator> {
1853 match &aggregate.function {
1854 AggregateFunction::Sum => Box::new(SumAccumulator::with_distinct_for_type(
1855 aggregate.distinct,
1856 aggregate.result_type.clone(),
1857 )),
1858 AggregateFunction::GroupConcat { separator } if !aggregate.order_by.is_empty() => {
1859 let sep = separator.clone().unwrap_or_else(|| ",".to_string());
1860 Box::new(GroupConcatAccumulator::with_order(
1861 sep,
1862 aggregate.distinct,
1863 aggregate_sort_specs(aggregate),
1864 ))
1865 }
1866 AggregateFunction::StringAgg { separator } if !aggregate.order_by.is_empty() => {
1867 let sep = separator.clone().unwrap_or_else(|| ",".to_string());
1868 Box::new(StringAggAccumulator::with_order(
1869 sep,
1870 aggregate.distinct,
1871 aggregate_sort_specs(aggregate),
1872 ))
1873 }
1874 AggregateFunction::PercentileDisc { fraction } => Box::new(PercentileDiscAccumulator::new(
1875 *fraction,
1876 aggregate_sort_specs(aggregate),
1877 )),
1878 AggregateFunction::PercentileCont { fraction }
1879 | AggregateFunction::QuantileCont { fraction } => Box::new(PercentileContAccumulator::new(
1880 *fraction,
1881 aggregate_sort_specs(aggregate)
1882 .first()
1883 .is_none_or(|(ascending, _)| *ascending),
1884 )),
1885 AggregateFunction::Mode if !aggregate.order_by.is_empty() => {
1886 Box::new(ModeAccumulator::new(aggregate_sort_specs(aggregate)))
1887 }
1888 AggregateFunction::First if !aggregate.order_by.is_empty() => Box::new(
1889 ValueAccumulator::new(ValueKind::First, aggregate_sort_specs(aggregate)),
1890 ),
1891 AggregateFunction::Last if !aggregate.order_by.is_empty() => Box::new(
1892 ValueAccumulator::new(ValueKind::Last, aggregate_sort_specs(aggregate)),
1893 ),
1894 _ => create_accumulator(&aggregate.function, aggregate.distinct),
1895 }
1896}
1897
1898pub fn exact_partial_aggregate_is_proven(aggregate: &AggregateExpr) -> bool {
1906 !aggregate.distinct
1907 && aggregate.order_by.is_empty()
1908 && matches!(
1909 aggregate.function,
1910 AggregateFunction::Count | AggregateFunction::Min | AggregateFunction::Max
1911 )
1912}
1913
1914pub fn merge_exact_aggregate_states(
1918 aggregates: &[AggregateExpr],
1919 partial_rows: impl IntoIterator<Item = Vec<Vec<SqlValue>>>,
1920) -> Result<Vec<SqlValue>> {
1921 if let Some(aggregate) = aggregates
1922 .iter()
1923 .find(|aggregate| !exact_partial_aggregate_is_proven(aggregate))
1924 {
1925 return Err(ExecutorError::InvalidOperation {
1926 operation: "distributed aggregate merge".into(),
1927 reason: format!(
1928 "{:?} requires ordered input replay rather than an exact partial merge",
1929 aggregate.function
1930 ),
1931 });
1932 }
1933
1934 let mut accumulators = aggregates
1935 .iter()
1936 .map(create_accumulator_for_aggregate)
1937 .collect::<Vec<_>>();
1938 for states in partial_rows {
1939 if states.len() != accumulators.len() {
1940 return Err(ExecutorError::InvalidOperation {
1941 operation: "distributed aggregate merge".into(),
1942 reason: format!(
1943 "partial state has {} aggregate(s), expected {}",
1944 states.len(),
1945 accumulators.len()
1946 ),
1947 });
1948 }
1949 for (accumulator, state) in accumulators.iter_mut().zip(states) {
1950 accumulator.merge(&state)?;
1951 }
1952 }
1953 accumulators
1954 .iter()
1955 .map(|accumulator| accumulator.finalize())
1956 .collect()
1957}
1958
1959const DEFAULT_GROUP_LIMIT: usize = 1_000_000;
1960const AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES: u64 = 32;
1961
1962#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1964pub enum AggregateMode {
1965 Partial,
1967 Final,
1969 Single,
1971}
1972
1973struct AggregateGroup {
1974 key_values: Vec<SqlValue>,
1975 accumulators: Vec<Box<dyn Accumulator>>,
1976 grouping_id: Option<i64>,
1979}
1980
1981pub struct AggregateIterator<'a> {
1983 input: Box<dyn RowIterator + 'a>,
1984 group_keys: Vec<TypedExpr>,
1985 aggregates: Vec<AggregateExpr>,
1986 having: Option<TypedExpr>,
1987 mode: AggregateMode,
1988 hash_table: Option<HashMap<GroupKeyBytes, AggregateGroup>>,
1989 result_rows: Vec<Row>,
1990 index: usize,
1991 schema: Vec<ColumnMetadata>,
1992 group_limit: usize,
1993 memory_tracker: Option<MemoryTracker>,
1994 shared_group_counter: Option<Arc<AtomicUsize>>,
1995 grouping_sets: Option<Vec<u64>>,
1999}
2000
2001impl<'a> AggregateIterator<'a> {
2002 pub fn new(
2004 input: Box<dyn RowIterator + 'a>,
2005 group_keys: Vec<TypedExpr>,
2006 aggregates: Vec<AggregateExpr>,
2007 having: Option<TypedExpr>,
2008 schema: Vec<ColumnMetadata>,
2009 ) -> Self {
2010 Self {
2011 input,
2012 group_keys,
2013 aggregates,
2014 having,
2015 mode: AggregateMode::Single,
2016 hash_table: None,
2017 result_rows: Vec::new(),
2018 index: 0,
2019 schema,
2020 group_limit: DEFAULT_GROUP_LIMIT,
2021 memory_tracker: None,
2022 shared_group_counter: None,
2023 grouping_sets: None,
2024 }
2025 }
2026
2027 pub fn with_group_limit(mut self, limit: usize) -> Self {
2029 self.group_limit = limit;
2030 self
2031 }
2032
2033 pub fn with_grouping_sets(mut self, grouping_sets: Option<Vec<u64>>) -> Self {
2041 self.grouping_sets = grouping_sets;
2042 self
2043 }
2044
2045 pub fn with_mode(mut self, mode: AggregateMode) -> Self {
2047 self.mode = mode;
2048 self
2049 }
2050
2051 pub fn with_memory_policy(mut self, policy: Option<MemoryPolicy>) -> Self {
2053 self.memory_tracker = policy.map(MemoryTracker::new);
2054 self
2055 }
2056
2057 pub fn with_shared_group_counter(mut self, counter: Option<Arc<AtomicUsize>>) -> Self {
2059 self.shared_group_counter = counter;
2060 self
2061 }
2062
2063 fn build_hash_table(&mut self) -> Result<()> {
2064 let mut table: HashMap<GroupKeyBytes, AggregateGroup> = HashMap::new();
2065 let mut next_row_id = 0u64;
2066
2067 while let Some(result) = self.input.next_row() {
2068 let row = result?;
2069 let (key_values, key_bytes) = match self.mode {
2070 AggregateMode::Final => {
2071 let key_values = row
2072 .values
2073 .get(..self.group_keys.len())
2074 .ok_or_else(|| {
2075 invalid_aggregate_state(
2076 "aggregate",
2077 "partial state row is missing group key values",
2078 )
2079 })?
2080 .to_vec();
2081 let key_bytes = encode_group_key(&key_values)?;
2082 (key_values, key_bytes)
2083 }
2084 AggregateMode::Partial | AggregateMode::Single => {
2085 let ctx = EvalContext::new(&row.values);
2086 let mut key_values = Vec::with_capacity(self.group_keys.len());
2087 for expr in &self.group_keys {
2088 key_values.push(crate::executor::evaluator::evaluate(expr, &ctx)?);
2089 }
2090 let key_bytes = encode_group_key(&key_values)?;
2091 (key_values, key_bytes)
2092 }
2093 };
2094
2095 let variants: Vec<(Vec<SqlValue>, GroupKeyBytes, Option<i64>)> =
2100 if let Some(sets) = &self.grouping_sets {
2101 let key_count = self.group_keys.len();
2102 let mut variants = Vec::with_capacity(sets.len());
2103 for (set_id, mask) in sets.iter().enumerate() {
2104 let mut masked = key_values.clone();
2105 for (position, value) in masked.iter_mut().enumerate() {
2106 if (mask >> (key_count - 1 - position)) & 1 == 1 {
2107 *value = SqlValue::Null;
2108 }
2109 }
2110 let mut bytes = Vec::with_capacity(4 + key_bytes.len());
2111 bytes.extend_from_slice(&(set_id as u32).to_le_bytes());
2112 bytes.extend_from_slice(&encode_group_key(&masked)?);
2113 variants.push((masked, bytes, Some(*mask as i64)));
2114 }
2115 variants
2116 } else {
2117 vec![(key_values, key_bytes, None)]
2118 };
2119
2120 for (key_values, key_bytes, grouping_id) in variants {
2121 if !table.contains_key(&key_bytes) {
2122 self.reserve_group_slot(table.len())?;
2123 if let Some(tracker) = &mut self.memory_tracker {
2124 tracker
2125 .add_values(&key_values)
2126 .map_err(map_core_memory_error)?;
2127 tracker
2128 .add_bytes(
2129 self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES,
2130 )
2131 .map_err(map_core_memory_error)?;
2132 }
2133 let accumulators = self
2134 .aggregates
2135 .iter()
2136 .map(|agg| {
2137 let mut aggregate = agg.clone();
2138 aggregate.distinct =
2139 matches!(self.mode, AggregateMode::Single) && agg.distinct;
2140 create_accumulator_for_aggregate(&aggregate)
2141 })
2142 .collect::<Vec<_>>();
2143 table.insert(
2144 key_bytes.clone(),
2145 AggregateGroup {
2146 key_values: key_values.clone(),
2147 accumulators,
2148 grouping_id,
2149 },
2150 );
2151 }
2152
2153 if let Some(group) = table.get_mut(&key_bytes) {
2154 match self.mode {
2155 AggregateMode::Final => {
2156 let mut offset = self.group_keys.len();
2157 for (idx, agg) in self.aggregates.iter().enumerate() {
2158 let arity = aggregate_state_types(agg).len();
2159 let state = row.values.get(offset..offset + arity).ok_or_else(|| {
2160 invalid_aggregate_state(
2161 "aggregate",
2162 format!(
2163 "partial state row is missing state values for aggregate {idx}"
2164 ),
2165 )
2166 })?;
2167 group.accumulators[idx].merge(state)?;
2168 offset += arity;
2169 }
2170 if offset != row.values.len() {
2171 return Err(invalid_aggregate_state(
2172 "aggregate",
2173 format!(
2174 "partial state row has {} trailing value(s)",
2175 row.values.len() - offset
2176 ),
2177 ));
2178 }
2179 }
2180 AggregateMode::Partial | AggregateMode::Single => {
2181 let ctx = EvalContext::new(&row.values);
2182 for (idx, agg) in self.aggregates.iter().enumerate() {
2183 if !aggregate_filter_accepts(agg, &ctx)? {
2187 continue;
2188 }
2189 let values = evaluate_aggregate_values(agg, &ctx)?;
2190 if let Some(tracker) = &mut self.memory_tracker
2191 && matches!(
2192 agg.function,
2193 AggregateFunction::GroupConcat { .. }
2194 | AggregateFunction::StringAgg { .. }
2195 | AggregateFunction::PercentileDisc { .. }
2196 | AggregateFunction::PercentileCont { .. }
2197 | AggregateFunction::QuantileCont { .. }
2198 | AggregateFunction::Median
2199 | AggregateFunction::Mode
2200 )
2201 && let Some(value_ref) = values.first()
2202 {
2203 tracker
2204 .add_value(value_ref)
2205 .map_err(map_core_memory_error)?;
2206 }
2207 if agg.order_by.is_empty() {
2208 group.accumulators[idx].update_values(&values)?;
2209 } else {
2210 let keys = evaluate_sort_keys(agg, &ctx)?;
2211 group.accumulators[idx]
2212 .update_ordered_values(&values, &keys)?;
2213 }
2214 }
2215 }
2216 }
2217 }
2218 }
2219 }
2220
2221 if let Some(sets) = self.grouping_sets.clone() {
2222 let key_count = self.group_keys.len();
2226 let full_mask = if key_count == 0 {
2227 0
2228 } else {
2229 (1u64 << key_count) - 1
2230 };
2231 if table.is_empty() {
2232 for (set_id, mask) in sets.iter().enumerate() {
2233 if *mask != full_mask {
2234 continue;
2235 }
2236 if let Some(tracker) = &mut self.memory_tracker {
2237 tracker
2238 .add_bytes(
2239 self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES,
2240 )
2241 .map_err(map_core_memory_error)?;
2242 }
2243 let accumulators = self
2244 .aggregates
2245 .iter()
2246 .map(|agg| {
2247 let mut aggregate = agg.clone();
2248 aggregate.distinct =
2249 matches!(self.mode, AggregateMode::Single) && agg.distinct;
2250 create_accumulator_for_aggregate(&aggregate)
2251 })
2252 .collect::<Vec<_>>();
2253 let key_values = vec![SqlValue::Null; key_count];
2254 let mut bytes = Vec::with_capacity(4);
2255 bytes.extend_from_slice(&(set_id as u32).to_le_bytes());
2256 bytes.extend_from_slice(&encode_group_key(&key_values)?);
2257 table.insert(
2258 bytes,
2259 AggregateGroup {
2260 key_values,
2261 accumulators,
2262 grouping_id: Some(*mask as i64),
2263 },
2264 );
2265 }
2266 }
2267 } else if table.is_empty() && self.group_keys.is_empty() {
2268 if let Some(tracker) = &mut self.memory_tracker {
2269 tracker
2270 .add_bytes(self.aggregates.len() as u64 * AGGREGATE_ACCUMULATOR_OVERHEAD_BYTES)
2271 .map_err(map_core_memory_error)?;
2272 }
2273 let accumulators = self
2274 .aggregates
2275 .iter()
2276 .map(|agg| {
2277 let mut aggregate = agg.clone();
2278 aggregate.distinct = matches!(self.mode, AggregateMode::Single) && agg.distinct;
2279 create_accumulator_for_aggregate(&aggregate)
2280 })
2281 .collect::<Vec<_>>();
2282 table.insert(
2283 Vec::new(),
2284 AggregateGroup {
2285 key_values: Vec::new(),
2286 accumulators,
2287 grouping_id: None,
2288 },
2289 );
2290 }
2291
2292 let mut rows = Vec::with_capacity(table.len());
2293 for group in table.values() {
2294 let mut values = Vec::with_capacity(self.group_keys.len() + self.aggregates.len() + 1);
2295 values.extend(group.key_values.iter().cloned());
2296 for acc in &group.accumulators {
2297 match self.mode {
2298 AggregateMode::Partial => values.extend(acc.state()?),
2299 AggregateMode::Final | AggregateMode::Single => values.push(acc.finalize()?),
2300 }
2301 }
2302 if let Some(grouping_id) = group.grouping_id {
2305 values.push(SqlValue::BigInt(grouping_id));
2306 }
2307 let row = Row::new(next_row_id, values);
2308 next_row_id += 1;
2309 if let Some(tracker) = &mut self.memory_tracker {
2310 tracker
2311 .add_row(&row.values)
2312 .map_err(map_core_memory_error)?;
2313 }
2314
2315 if self.mode != AggregateMode::Partial
2316 && let Some(having) = &self.having
2317 {
2318 let ctx = EvalContext::new(&row.values);
2319 match crate::executor::evaluator::evaluate(having, &ctx)? {
2320 SqlValue::Boolean(true) => rows.push(row),
2321 SqlValue::Boolean(false) | SqlValue::Null => {}
2322 other => {
2323 return Err(ExecutorError::Evaluation(
2324 crate::executor::EvaluationError::TypeMismatch {
2325 expected: "Boolean".into(),
2326 actual: other.type_name().into(),
2327 },
2328 ));
2329 }
2330 }
2331 } else {
2332 rows.push(row);
2333 }
2334 }
2335
2336 self.hash_table = Some(table);
2337 self.result_rows = rows;
2338 Ok(())
2339 }
2340
2341 fn reserve_group_slot(&self, local_group_count: usize) -> Result<()> {
2342 let next_count = if let Some(counter) = &self.shared_group_counter {
2343 counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1
2344 } else {
2345 local_group_count + 1
2346 };
2347 if next_count > self.group_limit {
2348 return Err(ExecutorError::ResourceExhausted {
2349 message: format!(
2350 "GROUP BY result exceeds memory limit (max groups: {})",
2351 self.group_limit
2352 ),
2353 });
2354 }
2355 Ok(())
2356 }
2357}
2358
2359impl<'a> RowIterator for AggregateIterator<'a> {
2360 fn next_row(&mut self) -> Option<Result<Row>> {
2361 if self.hash_table.is_none()
2362 && let Err(err) = self.build_hash_table()
2363 {
2364 return Some(Err(err));
2365 }
2366
2367 if self.index >= self.result_rows.len() {
2368 return None;
2369 }
2370 let row = self.result_rows[self.index].clone();
2371 self.index += 1;
2372 Some(Ok(row))
2373 }
2374
2375 fn schema(&self) -> &[ColumnMetadata] {
2376 &self.schema
2377 }
2378}
2379
2380pub struct StreamingAggregateIterator<'a> {
2382 input: Box<dyn RowIterator + 'a>,
2383 group_keys: Vec<TypedExpr>,
2384 aggregates: Vec<AggregateExpr>,
2385 having: Option<TypedExpr>,
2386 schema: Vec<ColumnMetadata>,
2387 current_key: Option<Vec<SqlValue>>,
2388 accumulators: Vec<Box<dyn Accumulator>>,
2389 pending_row: Option<Row>,
2390 finished: bool,
2391 next_row_id: u64,
2392 saw_row: bool,
2393}
2394
2395impl<'a> StreamingAggregateIterator<'a> {
2396 pub fn new(
2397 input: Box<dyn RowIterator + 'a>,
2398 group_keys: Vec<TypedExpr>,
2399 aggregates: Vec<AggregateExpr>,
2400 having: Option<TypedExpr>,
2401 schema: Vec<ColumnMetadata>,
2402 ) -> Self {
2403 Self {
2404 input,
2405 group_keys,
2406 aggregates,
2407 having,
2408 schema,
2409 current_key: None,
2410 accumulators: Vec::new(),
2411 pending_row: None,
2412 finished: false,
2413 next_row_id: 0,
2414 saw_row: false,
2415 }
2416 }
2417
2418 fn init_accumulators(&self) -> Vec<Box<dyn Accumulator>> {
2419 self.aggregates
2420 .iter()
2421 .map(create_accumulator_for_aggregate)
2422 .collect()
2423 }
2424
2425 fn update_accumulators(&mut self, ctx: &EvalContext<'_>) -> Result<()> {
2426 for (idx, agg) in self.aggregates.iter().enumerate() {
2427 if !aggregate_filter_accepts(agg, ctx)? {
2428 continue;
2429 }
2430 let values = evaluate_aggregate_values(agg, ctx)?;
2431 if agg.order_by.is_empty() {
2432 self.accumulators[idx].update_values(&values)?;
2433 } else {
2434 let keys = evaluate_sort_keys(agg, ctx)?;
2435 self.accumulators[idx].update_ordered_values(&values, &keys)?;
2436 }
2437 }
2438 Ok(())
2439 }
2440
2441 fn finalize_group(&mut self, key_values: &[SqlValue]) -> Result<Option<Row>> {
2442 let mut values = Vec::with_capacity(self.group_keys.len() + self.aggregates.len());
2443 values.extend(key_values.iter().cloned());
2444 for acc in &self.accumulators {
2445 values.push(acc.finalize()?);
2446 }
2447 let row = Row::new(self.next_row_id, values);
2448 self.next_row_id = self.next_row_id.saturating_add(1);
2449
2450 if let Some(having) = &self.having {
2451 let ctx = EvalContext::new(&row.values);
2452 match crate::executor::evaluator::evaluate(having, &ctx)? {
2453 SqlValue::Boolean(true) => Ok(Some(row)),
2454 SqlValue::Boolean(false) | SqlValue::Null => Ok(None),
2455 other => Err(ExecutorError::Evaluation(
2456 crate::executor::EvaluationError::TypeMismatch {
2457 expected: "Boolean".into(),
2458 actual: other.type_name().into(),
2459 },
2460 )),
2461 }
2462 } else {
2463 Ok(Some(row))
2464 }
2465 }
2466}
2467
2468impl<'a> RowIterator for StreamingAggregateIterator<'a> {
2469 fn next_row(&mut self) -> Option<Result<Row>> {
2470 if let Some(row) = self.pending_row.take() {
2471 return Some(Ok(row));
2472 }
2473 if self.finished {
2474 return None;
2475 }
2476
2477 loop {
2478 match self.input.next_row() {
2479 Some(Ok(row)) => {
2480 self.saw_row = true;
2481 let ctx = EvalContext::new(&row.values);
2482 let mut key_values = Vec::with_capacity(self.group_keys.len());
2483 for expr in &self.group_keys {
2484 match crate::executor::evaluator::evaluate(expr, &ctx) {
2485 Ok(value) => key_values.push(value),
2486 Err(err) => return Some(Err(err)),
2487 }
2488 }
2489
2490 match &self.current_key {
2491 None => {
2492 self.current_key = Some(key_values);
2493 self.accumulators = self.init_accumulators();
2494 if let Err(err) = self.update_accumulators(&ctx) {
2495 return Some(Err(err));
2496 }
2497 }
2498 Some(current_key) if *current_key == key_values => {
2499 if let Err(err) = self.update_accumulators(&ctx) {
2500 return Some(Err(err));
2501 }
2502 }
2503 Some(_) => {
2504 let current_key = self.current_key.clone().unwrap_or_default();
2505 let output = match self.finalize_group(¤t_key) {
2506 Ok(value) => value,
2507 Err(err) => return Some(Err(err)),
2508 };
2509 self.current_key = Some(key_values);
2510 self.accumulators = self.init_accumulators();
2511 if let Err(err) = self.update_accumulators(&ctx) {
2512 return Some(Err(err));
2513 }
2514 if let Some(row) = output {
2515 return Some(Ok(row));
2516 }
2517 }
2518 }
2519 }
2520 Some(Err(err)) => return Some(Err(err)),
2521 None => {
2522 self.finished = true;
2523 if let Some(current_key) = self.current_key.take() {
2524 return match self.finalize_group(¤t_key) {
2525 Ok(Some(row)) => Some(Ok(row)),
2526 Ok(None) => None,
2527 Err(err) => Some(Err(err)),
2528 };
2529 }
2530
2531 if self.group_keys.is_empty() && !self.saw_row {
2532 self.accumulators = self.init_accumulators();
2533 return match self.finalize_group(&[]) {
2534 Ok(Some(row)) => Some(Ok(row)),
2535 Ok(None) => None,
2536 Err(err) => Some(Err(err)),
2537 };
2538 }
2539
2540 return None;
2541 }
2542 }
2543 }
2544 }
2545
2546 fn schema(&self) -> &[ColumnMetadata] {
2547 &self.schema
2548 }
2549}
2550
2551fn aggregate_state_types(agg: &AggregateExpr) -> Vec<ResolvedType> {
2552 match &agg.function {
2553 AggregateFunction::Count => vec![ResolvedType::BigInt],
2554 AggregateFunction::Sum => vec![agg.result_type.clone()],
2555 AggregateFunction::Total => vec![ResolvedType::Double],
2556 AggregateFunction::Avg => vec![ResolvedType::Double, ResolvedType::BigInt],
2557 AggregateFunction::Min | AggregateFunction::Max => vec![agg.result_type.clone()],
2558 AggregateFunction::GroupConcat { .. } | AggregateFunction::StringAgg { .. } => {
2559 vec![ResolvedType::Text, ResolvedType::Text]
2560 }
2561 AggregateFunction::PercentileDisc { .. } => vec![agg.result_type.clone()],
2564 AggregateFunction::PercentileCont { .. }
2565 | AggregateFunction::QuantileCont { .. }
2566 | AggregateFunction::Median
2567 | AggregateFunction::Mode => vec![agg.result_type.clone()],
2568 AggregateFunction::Variance { .. }
2569 | AggregateFunction::Stddev { .. }
2570 | AggregateFunction::Covariance { .. }
2571 | AggregateFunction::Corr
2572 | AggregateFunction::RegrCount
2573 | AggregateFunction::RegrAvgX
2574 | AggregateFunction::RegrAvgY
2575 | AggregateFunction::RegrSxx
2576 | AggregateFunction::RegrSyy
2577 | AggregateFunction::RegrSxy
2578 | AggregateFunction::RegrSlope
2579 | AggregateFunction::RegrIntercept
2580 | AggregateFunction::RegrR2 => vec![
2581 ResolvedType::BigInt,
2582 ResolvedType::Double,
2583 ResolvedType::Double,
2584 ResolvedType::Double,
2585 ResolvedType::Double,
2586 ResolvedType::Double,
2587 ],
2588 AggregateFunction::AnyValue
2589 | AggregateFunction::First
2590 | AggregateFunction::Last
2591 | AggregateFunction::ArgMin
2592 | AggregateFunction::ArgMax
2593 | AggregateFunction::BitAnd
2594 | AggregateFunction::BitOr
2595 | AggregateFunction::BitXor
2596 | AggregateFunction::BoolAnd
2597 | AggregateFunction::BoolOr => vec![agg.result_type.clone()],
2598 }
2599}
2600
2601fn aggregate_filter_accepts(agg: &AggregateExpr, ctx: &EvalContext<'_>) -> Result<bool> {
2604 let Some(filter) = &agg.filter else {
2605 return Ok(true);
2606 };
2607 match crate::executor::evaluator::evaluate(filter, ctx)? {
2608 SqlValue::Boolean(true) => Ok(true),
2609 SqlValue::Boolean(false) | SqlValue::Null => Ok(false),
2610 other => Err(ExecutorError::Evaluation(
2611 crate::executor::EvaluationError::TypeMismatch {
2612 expected: "Boolean".into(),
2613 actual: other.type_name().into(),
2614 },
2615 )),
2616 }
2617}
2618
2619fn evaluate_sort_keys(agg: &AggregateExpr, ctx: &EvalContext<'_>) -> Result<Vec<SqlValue>> {
2621 agg.order_by
2622 .iter()
2623 .map(|sort| crate::executor::evaluator::evaluate(&sort.expr, ctx))
2624 .collect()
2625}
2626
2627fn evaluate_aggregate_values(agg: &AggregateExpr, ctx: &EvalContext<'_>) -> Result<Vec<SqlValue>> {
2628 agg.arg
2629 .iter()
2630 .chain(&agg.extra_args)
2631 .map(|expr| crate::executor::evaluator::evaluate(expr, ctx))
2632 .collect()
2633}
2634
2635pub fn build_aggregate_schema(
2637 group_keys: &[TypedExpr],
2638 aggregates: &[AggregateExpr],
2639) -> Vec<ColumnMetadata> {
2640 let mut schema = Vec::new();
2641 for (idx, key) in group_keys.iter().enumerate() {
2642 let name = match &key.kind {
2643 crate::planner::typed_expr::TypedExprKind::ColumnRef { column, .. } => column.clone(),
2644 _ => format!("group_{idx}"),
2645 };
2646 schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
2647 }
2648 for (idx, agg) in aggregates.iter().enumerate() {
2649 let name = match &agg.function {
2650 AggregateFunction::Count => format!("count_{idx}"),
2651 AggregateFunction::Sum => format!("sum_{idx}"),
2652 AggregateFunction::Total => format!("total_{idx}"),
2653 AggregateFunction::Avg => format!("avg_{idx}"),
2654 AggregateFunction::Min => format!("min_{idx}"),
2655 AggregateFunction::Max => format!("max_{idx}"),
2656 AggregateFunction::GroupConcat { .. } => format!("group_concat_{idx}"),
2657 AggregateFunction::StringAgg { .. } => format!("string_agg_{idx}"),
2658 AggregateFunction::PercentileDisc { .. } => format!("percentile_disc_{idx}"),
2659 AggregateFunction::PercentileCont { .. } => format!("percentile_cont_{idx}"),
2660 AggregateFunction::QuantileCont { .. } => format!("quantile_cont_{idx}"),
2661 AggregateFunction::Variance { sample } => {
2662 format!("var_{}_{idx}", if *sample { "samp" } else { "pop" })
2663 }
2664 AggregateFunction::Stddev { sample } => {
2665 format!("stddev_{}_{idx}", if *sample { "samp" } else { "pop" })
2666 }
2667 AggregateFunction::Covariance { sample } => {
2668 format!("covar_{}_{idx}", if *sample { "samp" } else { "pop" })
2669 }
2670 AggregateFunction::Corr => format!("corr_{idx}"),
2671 AggregateFunction::Median => format!("median_{idx}"),
2672 AggregateFunction::Mode => format!("mode_{idx}"),
2673 AggregateFunction::RegrCount => format!("regr_count_{idx}"),
2674 AggregateFunction::RegrAvgX => format!("regr_avgx_{idx}"),
2675 AggregateFunction::RegrAvgY => format!("regr_avgy_{idx}"),
2676 AggregateFunction::RegrSxx => format!("regr_sxx_{idx}"),
2677 AggregateFunction::RegrSyy => format!("regr_syy_{idx}"),
2678 AggregateFunction::RegrSxy => format!("regr_sxy_{idx}"),
2679 AggregateFunction::RegrSlope => format!("regr_slope_{idx}"),
2680 AggregateFunction::RegrIntercept => format!("regr_intercept_{idx}"),
2681 AggregateFunction::RegrR2 => format!("regr_r2_{idx}"),
2682 AggregateFunction::AnyValue => format!("any_value_{idx}"),
2683 AggregateFunction::First => format!("first_{idx}"),
2684 AggregateFunction::Last => format!("last_{idx}"),
2685 AggregateFunction::ArgMin => format!("arg_min_{idx}"),
2686 AggregateFunction::ArgMax => format!("arg_max_{idx}"),
2687 AggregateFunction::BitAnd => format!("bit_and_{idx}"),
2688 AggregateFunction::BitOr => format!("bit_or_{idx}"),
2689 AggregateFunction::BitXor => format!("bit_xor_{idx}"),
2690 AggregateFunction::BoolAnd => format!("bool_and_{idx}"),
2691 AggregateFunction::BoolOr => format!("bool_or_{idx}"),
2692 };
2693 schema.push(ColumnMetadata::new(name, agg.result_type.clone()));
2694 }
2695 schema
2696}
2697
2698pub fn build_partial_aggregate_schema(
2700 group_keys: &[TypedExpr],
2701 aggregates: &[AggregateExpr],
2702) -> Vec<ColumnMetadata> {
2703 let mut schema = Vec::new();
2704 for (idx, key) in group_keys.iter().enumerate() {
2705 let name = match &key.kind {
2706 crate::planner::typed_expr::TypedExprKind::ColumnRef { column, .. } => column.clone(),
2707 _ => format!("group_{idx}"),
2708 };
2709 schema.push(ColumnMetadata::new(name, key.resolved_type.clone()));
2710 }
2711 for (agg_idx, agg) in aggregates.iter().enumerate() {
2712 for (state_idx, state_type) in aggregate_state_types(agg).into_iter().enumerate() {
2713 schema.push(ColumnMetadata::new(
2714 format!("__agg{agg_idx}_state{state_idx}"),
2715 state_type,
2716 ));
2717 }
2718 }
2719 schema
2720}
2721
2722pub fn should_use_single_for_parallel(parallelism: usize, aggregates: &[AggregateExpr]) -> bool {
2728 parallelism <= 1
2729 || aggregates.iter().any(|agg| {
2730 agg.distinct
2731 || !agg.order_by.is_empty()
2732 || matches!(agg.function, AggregateFunction::PercentileDisc { .. })
2733 || matches!(
2734 agg.function,
2735 AggregateFunction::PercentileCont { .. }
2736 | AggregateFunction::QuantileCont { .. }
2737 | AggregateFunction::Median
2738 | AggregateFunction::Mode
2739 | AggregateFunction::AnyValue
2740 | AggregateFunction::First
2741 | AggregateFunction::Last
2742 | AggregateFunction::ArgMin
2743 | AggregateFunction::ArgMax
2744 )
2745 })
2746}
2747
2748fn collect_iterator_rows(iter: &mut dyn RowIterator) -> Result<Vec<Row>> {
2749 let mut rows = Vec::new();
2750 while let Some(result) = iter.next_row() {
2751 rows.push(result?);
2752 }
2753 Ok(rows)
2754}
2755
2756struct ChainRowIterator<'a> {
2757 prefix: std::vec::IntoIter<Row>,
2758 tail: Box<dyn RowIterator + 'a>,
2759 schema: Vec<ColumnMetadata>,
2760}
2761
2762impl<'a> ChainRowIterator<'a> {
2763 fn new(prefix: Vec<Row>, tail: Box<dyn RowIterator + 'a>, schema: Vec<ColumnMetadata>) -> Self {
2764 Self {
2765 prefix: prefix.into_iter(),
2766 tail,
2767 schema,
2768 }
2769 }
2770}
2771
2772impl RowIterator for ChainRowIterator<'_> {
2773 fn next_row(&mut self) -> Option<Result<Row>> {
2774 if let Some(row) = self.prefix.next() {
2775 return Some(Ok(row));
2776 }
2777 self.tail.next_row()
2778 }
2779
2780 fn schema(&self) -> &[ColumnMetadata] {
2781 &self.schema
2782 }
2783}
2784
2785fn estimate_row_bytes(row: &Row) -> u64 {
2786 row.values.iter().map(ByteSized::estimated_bytes).sum()
2787}
2788
2789fn split_contiguous_partitions(rows: Vec<Row>, parallelism: usize) -> Vec<Vec<Row>> {
2790 let requested = parallelism.max(1);
2791 if rows.is_empty() {
2792 return (0..requested).map(|_| Vec::new()).collect();
2793 }
2794 let partitions = requested.min(rows.len());
2795 let total = rows.len();
2796 let mut tail = rows;
2797 let mut output = Vec::with_capacity(partitions);
2798 for partition in (0..partitions).rev() {
2799 let start = partition * total / partitions;
2800 output.push(tail.split_off(start));
2801 }
2802 output.reverse();
2803 output
2804}
2805
2806#[allow(clippy::too_many_arguments)]
2807fn execute_partial_partition(
2808 partition_index: usize,
2809 rows: Vec<Row>,
2810 input_schema: Vec<ColumnMetadata>,
2811 group_keys: Vec<TypedExpr>,
2812 aggregates: Vec<AggregateExpr>,
2813 partial_schema: Vec<ColumnMetadata>,
2814 group_limit: usize,
2815 shared_group_counter: Arc<AtomicUsize>,
2816 shared_memory_counter: Arc<AtomicU64>,
2817 memory_limit: Option<u64>,
2818) -> Result<(usize, Vec<Row>)> {
2819 let input = VecIterator::new(rows, input_schema);
2820 let mut iter = AggregateIterator::new(
2821 Box::new(input),
2822 group_keys,
2823 aggregates,
2824 None,
2825 partial_schema,
2826 )
2827 .with_mode(AggregateMode::Partial)
2828 .with_group_limit(group_limit)
2829 .with_shared_group_counter(Some(shared_group_counter));
2830 let rows = collect_iterator_rows(&mut iter)?;
2831 for row in &rows {
2832 let used = shared_memory_counter
2833 .fetch_add(estimate_row_bytes(row), std::sync::atomic::Ordering::SeqCst)
2834 .saturating_add(estimate_row_bytes(row));
2835 if let Some(limit) = memory_limit
2836 && used > limit
2837 {
2838 return Err(ExecutorError::ResourceExhausted {
2839 message: format!(
2840 "parallel aggregate memory limit exceeded: {used} bytes (limit {limit})"
2841 ),
2842 });
2843 }
2844 }
2845 Ok((partition_index, rows))
2846}
2847
2848fn recv_partition_results(
2849 receiver: std::sync::mpsc::Receiver<Result<(usize, Vec<Row>)>>,
2850 expected: usize,
2851) -> Result<Vec<(usize, Vec<Row>)>> {
2852 let mut outputs = Vec::with_capacity(expected);
2853 for _ in 0..expected {
2854 let result = receiver
2855 .recv()
2856 .map_err(|err| ExecutorError::InvalidOperation {
2857 operation: "parallel aggregate".into(),
2858 reason: format!("partition worker failed to report result: {err}"),
2859 })?;
2860 outputs.push(result?);
2861 }
2862 outputs.sort_by_key(|(idx, _)| *idx);
2863 Ok(outputs)
2864}
2865
2866#[cfg(feature = "tokio")]
2867#[allow(clippy::too_many_arguments)]
2868fn run_partial_partitions(
2869 partitions: Vec<Vec<Row>>,
2870 input_schema: Vec<ColumnMetadata>,
2871 group_keys: Vec<TypedExpr>,
2872 aggregates: Vec<AggregateExpr>,
2873 partial_schema: Vec<ColumnMetadata>,
2874 group_limit: usize,
2875 shared_group_counter: Arc<AtomicUsize>,
2876 shared_memory_counter: Arc<AtomicU64>,
2877 memory_limit: Option<u64>,
2878) -> Result<Vec<(usize, Vec<Row>)>> {
2879 if let Ok(handle) = tokio::runtime::Handle::try_current() {
2880 let expected = partitions.len();
2881 let (sender, receiver) = std::sync::mpsc::channel();
2882 for (partition_index, rows) in partitions.into_iter().enumerate() {
2883 let sender = sender.clone();
2884 let input_schema = input_schema.clone();
2885 let group_keys = group_keys.clone();
2886 let aggregates = aggregates.clone();
2887 let partial_schema = partial_schema.clone();
2888 let shared_group_counter = Arc::clone(&shared_group_counter);
2889 let shared_memory_counter = Arc::clone(&shared_memory_counter);
2890 handle.spawn_blocking(move || {
2891 let result = execute_partial_partition(
2892 partition_index,
2893 rows,
2894 input_schema,
2895 group_keys,
2896 aggregates,
2897 partial_schema,
2898 group_limit,
2899 shared_group_counter,
2900 shared_memory_counter,
2901 memory_limit,
2902 );
2903 let _ = sender.send(result);
2904 });
2905 }
2906 drop(sender);
2907 return recv_partition_results(receiver, expected);
2908 }
2909
2910 run_partial_partitions_on_threads(
2911 partitions,
2912 input_schema,
2913 group_keys,
2914 aggregates,
2915 partial_schema,
2916 group_limit,
2917 shared_group_counter,
2918 shared_memory_counter,
2919 memory_limit,
2920 )
2921}
2922
2923#[cfg(not(feature = "tokio"))]
2924#[allow(clippy::too_many_arguments)]
2925fn run_partial_partitions(
2926 partitions: Vec<Vec<Row>>,
2927 input_schema: Vec<ColumnMetadata>,
2928 group_keys: Vec<TypedExpr>,
2929 aggregates: Vec<AggregateExpr>,
2930 partial_schema: Vec<ColumnMetadata>,
2931 group_limit: usize,
2932 shared_group_counter: Arc<AtomicUsize>,
2933 shared_memory_counter: Arc<AtomicU64>,
2934 memory_limit: Option<u64>,
2935) -> Result<Vec<(usize, Vec<Row>)>> {
2936 run_partial_partitions_on_threads(
2937 partitions,
2938 input_schema,
2939 group_keys,
2940 aggregates,
2941 partial_schema,
2942 group_limit,
2943 shared_group_counter,
2944 shared_memory_counter,
2945 memory_limit,
2946 )
2947}
2948
2949#[allow(clippy::too_many_arguments)]
2950fn run_partial_partitions_on_threads(
2951 partitions: Vec<Vec<Row>>,
2952 input_schema: Vec<ColumnMetadata>,
2953 group_keys: Vec<TypedExpr>,
2954 aggregates: Vec<AggregateExpr>,
2955 partial_schema: Vec<ColumnMetadata>,
2956 group_limit: usize,
2957 shared_group_counter: Arc<AtomicUsize>,
2958 shared_memory_counter: Arc<AtomicU64>,
2959 memory_limit: Option<u64>,
2960) -> Result<Vec<(usize, Vec<Row>)>> {
2961 let expected = partitions.len();
2962 let (sender, receiver) = std::sync::mpsc::channel();
2963 std::thread::scope(|scope| {
2964 for (partition_index, rows) in partitions.into_iter().enumerate() {
2965 let sender = sender.clone();
2966 let input_schema = input_schema.clone();
2967 let group_keys = group_keys.clone();
2968 let aggregates = aggregates.clone();
2969 let partial_schema = partial_schema.clone();
2970 let shared_group_counter = Arc::clone(&shared_group_counter);
2971 let shared_memory_counter = Arc::clone(&shared_memory_counter);
2972 scope.spawn(move || {
2973 let result = execute_partial_partition(
2974 partition_index,
2975 rows,
2976 input_schema,
2977 group_keys,
2978 aggregates,
2979 partial_schema,
2980 group_limit,
2981 shared_group_counter,
2982 shared_memory_counter,
2983 memory_limit,
2984 );
2985 let _ = sender.send(result);
2986 });
2987 }
2988 });
2989 drop(sender);
2990 recv_partition_results(receiver, expected)
2991}
2992
2993pub fn execute_parallel_aggregate_rows<'a>(
2995 input: Box<dyn RowIterator + 'a>,
2996 group_keys: Vec<TypedExpr>,
2997 aggregates: Vec<AggregateExpr>,
2998 having: Option<TypedExpr>,
2999 final_schema: Vec<ColumnMetadata>,
3000 parallelism: usize,
3001) -> Result<Vec<Row>> {
3002 execute_parallel_aggregate_rows_with_policy(
3003 input,
3004 group_keys,
3005 aggregates,
3006 having,
3007 final_schema,
3008 parallelism,
3009 None,
3010 DEFAULT_GROUP_LIMIT,
3011 )
3012}
3013
3014#[allow(clippy::too_many_arguments)]
3016pub fn execute_parallel_aggregate_rows_with_policy<'a>(
3017 mut input: Box<dyn RowIterator + 'a>,
3018 group_keys: Vec<TypedExpr>,
3019 aggregates: Vec<AggregateExpr>,
3020 having: Option<TypedExpr>,
3021 final_schema: Vec<ColumnMetadata>,
3022 parallelism: usize,
3023 memory: Option<MemoryPolicy>,
3024 group_limit: usize,
3025) -> Result<Vec<Row>> {
3026 if parallelism <= 1 {
3027 return execute_single_aggregate_rows(
3028 input,
3029 group_keys,
3030 aggregates,
3031 having,
3032 final_schema,
3033 memory,
3034 group_limit,
3035 );
3036 }
3037
3038 let input_schema = input.schema().to_vec();
3039 let mut input_rows = Vec::new();
3040 let mut materialized_bytes = 0u64;
3041 let materialize_threshold = memory
3042 .as_ref()
3043 .and_then(MemoryPolicy::limit_bytes)
3044 .map(|limit| (limit / 2).max(1));
3045
3046 while let Some(result) = input.next_row() {
3047 let row = result?;
3048 let row_bytes = materialize_threshold.map(|_| estimate_row_bytes(&row));
3049 if let (Some(threshold), Some(row_bytes)) = (materialize_threshold, row_bytes)
3050 && materialized_bytes.saturating_add(row_bytes) > threshold
3051 {
3052 input_rows.push(row);
3053 let chained = ChainRowIterator::new(input_rows, input, input_schema.clone());
3054 return execute_single_aggregate_rows(
3055 Box::new(chained),
3056 group_keys,
3057 aggregates,
3058 having,
3059 final_schema,
3060 memory,
3061 group_limit,
3062 );
3063 }
3064 if let Some(row_bytes) = row_bytes {
3065 materialized_bytes = materialized_bytes.saturating_add(row_bytes);
3066 }
3067 input_rows.push(row);
3068 }
3069
3070 let fallback_rows = if memory.is_some() || group_limit < input_rows.len() {
3071 Some(input_rows.clone())
3072 } else {
3073 None
3074 };
3075 let result = execute_parallel_aggregate_rows_from_materialized(
3076 input_rows,
3077 input_schema.clone(),
3078 group_keys.clone(),
3079 aggregates.clone(),
3080 having.clone(),
3081 final_schema.clone(),
3082 parallelism,
3083 group_limit,
3084 materialized_bytes,
3085 memory.as_ref().and_then(MemoryPolicy::limit_bytes),
3086 );
3087 match result {
3088 Ok(rows) => Ok(rows),
3089 Err(ExecutorError::ResourceExhausted { .. }) => {
3090 if let Some(fallback_rows) = fallback_rows {
3091 execute_single_aggregate_rows(
3092 Box::new(VecIterator::new(fallback_rows, input_schema)),
3093 group_keys,
3094 aggregates,
3095 having,
3096 final_schema,
3097 memory,
3098 group_limit,
3099 )
3100 } else {
3101 Err(ExecutorError::ResourceExhausted {
3102 message: format!(
3103 "parallel aggregate exceeded group limit {group_limit}; no fallback rows retained"
3104 ),
3105 })
3106 }
3107 }
3108 Err(err) => Err(err),
3109 }
3110}
3111
3112#[allow(clippy::too_many_arguments)]
3113fn execute_parallel_aggregate_rows_from_materialized(
3114 input_rows: Vec<Row>,
3115 input_schema: Vec<ColumnMetadata>,
3116 group_keys: Vec<TypedExpr>,
3117 aggregates: Vec<AggregateExpr>,
3118 having: Option<TypedExpr>,
3119 final_schema: Vec<ColumnMetadata>,
3120 parallelism: usize,
3121 group_limit: usize,
3122 materialized_bytes: u64,
3123 memory_limit: Option<u64>,
3124) -> Result<Vec<Row>> {
3125 let partial_schema = build_partial_aggregate_schema(&group_keys, &aggregates);
3126 let partitions = split_contiguous_partitions(input_rows, parallelism);
3127 let shared_group_counter = Arc::new(AtomicUsize::new(0));
3128 let shared_memory_counter = Arc::new(AtomicU64::new(materialized_bytes));
3129 let partial_results = run_partial_partitions(
3130 partitions,
3131 input_schema,
3132 group_keys.clone(),
3133 aggregates.clone(),
3134 partial_schema.clone(),
3135 group_limit,
3136 shared_group_counter,
3137 shared_memory_counter,
3138 memory_limit,
3139 )?;
3140
3141 let partial_rows = partial_results
3142 .into_iter()
3143 .flat_map(|(_, rows)| rows)
3144 .collect::<Vec<_>>();
3145 let final_input = VecIterator::new(partial_rows, partial_schema);
3146 let mut final_iter = AggregateIterator::new(
3147 Box::new(final_input),
3148 group_keys,
3149 aggregates,
3150 having,
3151 final_schema,
3152 )
3153 .with_mode(AggregateMode::Final)
3154 .with_group_limit(group_limit);
3155 collect_iterator_rows(&mut final_iter)
3156}
3157
3158fn execute_single_aggregate_rows<'a>(
3159 input: Box<dyn RowIterator + 'a>,
3160 group_keys: Vec<TypedExpr>,
3161 aggregates: Vec<AggregateExpr>,
3162 having: Option<TypedExpr>,
3163 final_schema: Vec<ColumnMetadata>,
3164 memory: Option<MemoryPolicy>,
3165 group_limit: usize,
3166) -> Result<Vec<Row>> {
3167 let mut iter = AggregateIterator::new(input, group_keys, aggregates, having, final_schema)
3168 .with_group_limit(group_limit)
3169 .with_memory_policy(memory);
3170 collect_iterator_rows(&mut iter)
3171}
3172
3173#[cfg(test)]
3174mod tests {
3175 use super::*;
3176 use crate::ast::span::Span;
3177 use crate::executor::memory::SpillPolicy;
3178 use crate::planner::typed_expr::TypedExprKind;
3179
3180 fn apply_values(acc: &mut dyn Accumulator, values: &[Option<SqlValue>]) {
3181 for value in values {
3182 acc.update(value.clone()).unwrap();
3183 }
3184 }
3185
3186 fn single_result(
3187 make_accumulator: impl Fn() -> Box<dyn Accumulator>,
3188 partitions: &[Vec<Option<SqlValue>>],
3189 ) -> SqlValue {
3190 let mut acc = make_accumulator();
3191 for partition in partitions {
3192 apply_values(acc.as_mut(), partition);
3193 }
3194 acc.finalize().unwrap()
3195 }
3196
3197 fn merged_result(
3198 make_partial: impl Fn() -> Box<dyn Accumulator>,
3199 make_final: impl Fn() -> Box<dyn Accumulator>,
3200 partitions: &[Vec<Option<SqlValue>>],
3201 merge_order: &[usize],
3202 ) -> SqlValue {
3203 let states = partitions
3204 .iter()
3205 .map(|partition| {
3206 let mut acc = make_partial();
3207 apply_values(acc.as_mut(), partition);
3208 acc.state().unwrap()
3209 })
3210 .collect::<Vec<_>>();
3211
3212 let mut final_acc = make_final();
3213 for idx in merge_order {
3214 final_acc.merge(&states[*idx]).unwrap();
3215 }
3216 final_acc.finalize().unwrap()
3217 }
3218
3219 fn assert_single_equals_merged(
3220 make_accumulator: impl Fn() -> Box<dyn Accumulator> + Copy,
3221 partitions: Vec<Vec<Option<SqlValue>>>,
3222 ) {
3223 let merge_order = (0..partitions.len()).collect::<Vec<_>>();
3224 let single = single_result(make_accumulator, &partitions);
3225 let merged = merged_result(
3226 make_accumulator,
3227 make_accumulator,
3228 &partitions,
3229 &merge_order,
3230 );
3231 assert_eq!(single, merged);
3232 }
3233
3234 fn assert_merge_order_invariant(
3235 make_accumulator: impl Fn() -> Box<dyn Accumulator> + Copy,
3236 partitions: Vec<Vec<Option<SqlValue>>>,
3237 merge_orders: &[Vec<usize>],
3238 ) {
3239 let single = single_result(make_accumulator, &partitions);
3240 for order in merge_orders {
3241 let merged = merged_result(make_accumulator, make_accumulator, &partitions, order);
3242 assert_eq!(single, merged, "merge order {order:?}");
3243 }
3244 }
3245
3246 fn column_expr(index: usize, name: &str, resolved_type: ResolvedType) -> TypedExpr {
3247 TypedExpr {
3248 kind: TypedExprKind::ColumnRef {
3249 table: "t".into(),
3250 column: name.into(),
3251 column_index: index,
3252 },
3253 resolved_type,
3254 span: Span::default(),
3255 }
3256 }
3257
3258 fn sample_aggregate_schema() -> Vec<ColumnMetadata> {
3259 vec![
3260 ColumnMetadata::new("category", ResolvedType::Text),
3261 ColumnMetadata::new("price", ResolvedType::Double),
3262 ColumnMetadata::new("label", ResolvedType::Text),
3263 ]
3264 }
3265
3266 fn sample_aggregate_rows() -> Vec<Row> {
3267 vec![
3268 Row::new(
3269 0,
3270 vec![
3271 SqlValue::Text("book".into()),
3272 SqlValue::Double(10.0),
3273 SqlValue::Text("a".into()),
3274 ],
3275 ),
3276 Row::new(
3277 1,
3278 vec![
3279 SqlValue::Text("book".into()),
3280 SqlValue::Double(15.0),
3281 SqlValue::Text("b".into()),
3282 ],
3283 ),
3284 Row::new(
3285 2,
3286 vec![
3287 SqlValue::Text("game".into()),
3288 SqlValue::Double(20.0),
3289 SqlValue::Text("c".into()),
3290 ],
3291 ),
3292 Row::new(
3293 3,
3294 vec![
3295 SqlValue::Text("book".into()),
3296 SqlValue::Null,
3297 SqlValue::Text("a".into()),
3298 ],
3299 ),
3300 Row::new(
3301 4,
3302 vec![
3303 SqlValue::Text("toy".into()),
3304 SqlValue::Double(3.0),
3305 SqlValue::Text("d".into()),
3306 ],
3307 ),
3308 ]
3309 }
3310
3311 fn sample_aggregates() -> Vec<AggregateExpr> {
3312 let price = column_expr(1, "price", ResolvedType::Double);
3313 let label = column_expr(2, "label", ResolvedType::Text);
3314 vec![
3315 AggregateExpr::count_star(),
3316 AggregateExpr::sum(price.clone()),
3317 AggregateExpr::avg(price),
3318 AggregateExpr {
3319 function: AggregateFunction::GroupConcat {
3320 separator: Some("|".into()),
3321 },
3322 arg: Some(label),
3323 extra_args: Vec::new(),
3324 distinct: false,
3325 result_type: ResolvedType::Text,
3326 filter: None,
3327 order_by: Vec::new(),
3328 },
3329 ]
3330 }
3331
3332 fn collect_single_aggregate(
3333 group_keys: Vec<TypedExpr>,
3334 aggregates: Vec<AggregateExpr>,
3335 ) -> Vec<Vec<SqlValue>> {
3336 let input = VecIterator::new(sample_aggregate_rows(), sample_aggregate_schema());
3337 let schema = build_aggregate_schema(&group_keys, &aggregates);
3338 let mut iter =
3339 AggregateIterator::new(Box::new(input), group_keys, aggregates, None, schema);
3340 collect_iterator_rows(&mut iter)
3341 .unwrap()
3342 .into_iter()
3343 .map(|row| row.values)
3344 .collect()
3345 }
3346
3347 fn collect_parallel_aggregate(
3348 group_keys: Vec<TypedExpr>,
3349 aggregates: Vec<AggregateExpr>,
3350 parallelism: usize,
3351 ) -> Vec<Vec<SqlValue>> {
3352 let input = VecIterator::new(sample_aggregate_rows(), sample_aggregate_schema());
3353 let schema = build_aggregate_schema(&group_keys, &aggregates);
3354 execute_parallel_aggregate_rows(
3355 Box::new(input),
3356 group_keys,
3357 aggregates,
3358 None,
3359 schema,
3360 parallelism,
3361 )
3362 .unwrap()
3363 .into_iter()
3364 .map(|row| row.values)
3365 .collect()
3366 }
3367
3368 fn sort_rows(mut rows: Vec<Vec<SqlValue>>) -> Vec<Vec<SqlValue>> {
3369 rows.sort_by(|left, right| format!("{left:?}").cmp(&format!("{right:?}")));
3370 rows
3371 }
3372
3373 #[test]
3374 fn grouping_sets_accumulate_each_row_once_per_set() {
3375 let category = column_expr(0, "category", ResolvedType::Text);
3376 let group_keys = vec![category];
3377 let aggregates = vec![AggregateExpr::count_star()];
3378 let input = VecIterator::new(sample_aggregate_rows(), sample_aggregate_schema());
3379 let mut schema = build_aggregate_schema(&group_keys, &aggregates);
3380 schema.push(ColumnMetadata::new("__grouping_id", ResolvedType::BigInt));
3381 let mut iter =
3383 AggregateIterator::new(Box::new(input), group_keys, aggregates, None, schema)
3384 .with_grouping_sets(Some(vec![0b0, 0b1]));
3385 let rows = sort_rows(
3386 collect_iterator_rows(&mut iter)
3387 .unwrap()
3388 .into_iter()
3389 .map(|row| row.values)
3390 .collect(),
3391 );
3392
3393 assert_eq!(
3394 rows,
3395 sort_rows(vec![
3396 vec![
3397 SqlValue::Text("book".into()),
3398 SqlValue::BigInt(3),
3399 SqlValue::BigInt(0),
3400 ],
3401 vec![
3402 SqlValue::Text("game".into()),
3403 SqlValue::BigInt(1),
3404 SqlValue::BigInt(0),
3405 ],
3406 vec![
3407 SqlValue::Text("toy".into()),
3408 SqlValue::BigInt(1),
3409 SqlValue::BigInt(0),
3410 ],
3411 vec![SqlValue::Null, SqlValue::BigInt(5), SqlValue::BigInt(1)],
3412 ])
3413 );
3414 }
3415
3416 #[test]
3417 fn grouping_sets_group_limit_applies_across_all_sets() {
3418 let category = column_expr(0, "category", ResolvedType::Text);
3419 let group_keys = vec![category];
3420 let aggregates = vec![AggregateExpr::count_star()];
3421 let input = VecIterator::new(sample_aggregate_rows(), sample_aggregate_schema());
3422 let mut schema = build_aggregate_schema(&group_keys, &aggregates);
3423 schema.push(ColumnMetadata::new("__grouping_id", ResolvedType::BigInt));
3424 let mut iter =
3428 AggregateIterator::new(Box::new(input), group_keys, aggregates, None, schema)
3429 .with_grouping_sets(Some(vec![0b0, 0b1]))
3430 .with_group_limit(3);
3431 let error = collect_iterator_rows(&mut iter).unwrap_err();
3432 assert!(matches!(error, ExecutorError::ResourceExhausted { .. }));
3433 }
3434
3435 #[test]
3436 fn partial_schema_uses_group_keys_and_state_columns() {
3437 let category = column_expr(0, "category", ResolvedType::Text);
3438 let price = column_expr(1, "price", ResolvedType::Double);
3439 let aggregates = vec![AggregateExpr::count_star(), AggregateExpr::avg(price)];
3440
3441 let schema = build_partial_aggregate_schema(&[category], &aggregates);
3442 let names = schema
3443 .iter()
3444 .map(|column| column.name.as_str())
3445 .collect::<Vec<_>>();
3446 assert_eq!(
3447 names,
3448 vec![
3449 "category",
3450 "__agg0_state0",
3451 "__agg1_state0",
3452 "__agg1_state1"
3453 ]
3454 );
3455 assert_eq!(schema[1].data_type, ResolvedType::BigInt);
3456 assert_eq!(schema[2].data_type, ResolvedType::Double);
3457 assert_eq!(schema[3].data_type, ResolvedType::BigInt);
3458 }
3459
3460 #[test]
3461 fn parallel_aggregate_matches_single_with_group_by() {
3462 let group_keys = vec![column_expr(0, "category", ResolvedType::Text)];
3463 let aggregates = sample_aggregates();
3464
3465 let single = sort_rows(collect_single_aggregate(
3466 group_keys.clone(),
3467 aggregates.clone(),
3468 ));
3469 let parallel = sort_rows(collect_parallel_aggregate(group_keys, aggregates, 3));
3470
3471 assert_eq!(parallel, single);
3472 }
3473
3474 #[test]
3475 fn parallel_aggregate_matches_single_without_group_by() {
3476 let aggregates = sample_aggregates();
3477
3478 let single = collect_single_aggregate(Vec::new(), aggregates.clone());
3479 let parallel = collect_parallel_aggregate(Vec::new(), aggregates, 4);
3480
3481 assert_eq!(parallel, single);
3482 }
3483
3484 #[test]
3485 fn distinct_aggregates_force_single_parallel_mode() {
3486 let price = column_expr(1, "price", ResolvedType::Double);
3487 let aggregates = vec![AggregateExpr {
3488 distinct: true,
3489 ..AggregateExpr::sum(price)
3490 }];
3491
3492 assert!(should_use_single_for_parallel(4, &aggregates));
3493 assert!(should_use_single_for_parallel(1, &sample_aggregates()));
3494 assert!(!should_use_single_for_parallel(2, &sample_aggregates()));
3495 }
3496
3497 #[test]
3498 fn parallel_group_counter_exhaustion_falls_back_to_single() {
3499 let schema = vec![ColumnMetadata::new("category", ResolvedType::Text)];
3500 let rows = vec![
3501 Row::new(0, vec![SqlValue::Text("a".into())]),
3502 Row::new(1, vec![SqlValue::Text("b".into())]),
3503 Row::new(2, vec![SqlValue::Text("a".into())]),
3504 Row::new(3, vec![SqlValue::Text("b".into())]),
3505 ];
3506 let group_keys = vec![column_expr(0, "category", ResolvedType::Text)];
3507 let aggregates = vec![AggregateExpr::count_star()];
3508 let final_schema = build_aggregate_schema(&group_keys, &aggregates);
3509
3510 let single_values = {
3511 let input = VecIterator::new(rows.clone(), schema.clone());
3512 execute_single_aggregate_rows(
3513 Box::new(input),
3514 group_keys.clone(),
3515 aggregates.clone(),
3516 None,
3517 final_schema.clone(),
3518 None,
3519 2,
3520 )
3521 .unwrap()
3522 .into_iter()
3523 .map(|row| row.values)
3524 .collect::<Vec<_>>()
3525 };
3526
3527 let parallel_values = execute_parallel_aggregate_rows_with_policy(
3528 Box::new(VecIterator::new(rows, schema)),
3529 group_keys,
3530 aggregates,
3531 None,
3532 final_schema,
3533 2,
3534 None,
3535 2,
3536 )
3537 .unwrap()
3538 .into_iter()
3539 .map(|row| row.values)
3540 .collect::<Vec<_>>();
3541
3542 assert_eq!(sort_rows(parallel_values), sort_rows(single_values));
3543 }
3544
3545 #[test]
3546 fn materialize_limit_exhaustion_falls_back_to_streaming_single() {
3547 let schema = vec![ColumnMetadata::new("payload", ResolvedType::Text)];
3548 let rows = (0..4)
3549 .map(|idx| Row::new(idx, vec![SqlValue::Text("x".repeat(40))]))
3550 .collect::<Vec<_>>();
3551 let aggregates = vec![AggregateExpr::count_star()];
3552 let final_schema = build_aggregate_schema(&[], &aggregates);
3553 let policy = MemoryPolicy::new(Some(100), SpillPolicy::FailFast);
3554
3555 let result = execute_parallel_aggregate_rows_with_policy(
3556 Box::new(VecIterator::new(rows, schema)),
3557 Vec::new(),
3558 aggregates,
3559 None,
3560 final_schema,
3561 4,
3562 Some(policy),
3563 DEFAULT_GROUP_LIMIT,
3564 )
3565 .unwrap();
3566
3567 assert_eq!(result.len(), 1);
3568 assert_eq!(result[0].values, vec![SqlValue::BigInt(4)]);
3569 }
3570
3571 #[test]
3572 fn streaming_aggregate_respects_distinct_accumulators() {
3573 let schema = sample_aggregate_schema();
3574 let rows = vec![
3575 Row::new(
3576 0,
3577 vec![
3578 SqlValue::Text("book".into()),
3579 SqlValue::Double(10.0),
3580 SqlValue::Text("a".into()),
3581 ],
3582 ),
3583 Row::new(
3584 1,
3585 vec![
3586 SqlValue::Text("book".into()),
3587 SqlValue::Double(10.0),
3588 SqlValue::Text("a".into()),
3589 ],
3590 ),
3591 Row::new(
3592 2,
3593 vec![
3594 SqlValue::Text("book".into()),
3595 SqlValue::Double(15.0),
3596 SqlValue::Text("b".into()),
3597 ],
3598 ),
3599 ];
3600 let group_keys = vec![column_expr(0, "category", ResolvedType::Text)];
3601 let price = column_expr(1, "price", ResolvedType::Double);
3602 let label = column_expr(2, "label", ResolvedType::Text);
3603 let aggregates = vec![
3604 AggregateExpr {
3605 distinct: true,
3606 ..AggregateExpr::sum(price)
3607 },
3608 AggregateExpr {
3609 function: AggregateFunction::GroupConcat {
3610 separator: Some("|".into()),
3611 },
3612 arg: Some(label),
3613 extra_args: Vec::new(),
3614 distinct: true,
3615 result_type: ResolvedType::Text,
3616 filter: None,
3617 order_by: Vec::new(),
3618 },
3619 ];
3620 let output_schema = build_aggregate_schema(&group_keys, &aggregates);
3621 let input = VecIterator::new(rows, schema);
3622 let mut iter = StreamingAggregateIterator::new(
3623 Box::new(input),
3624 group_keys,
3625 aggregates,
3626 None,
3627 output_schema,
3628 );
3629 let rows = collect_iterator_rows(&mut iter).unwrap();
3630
3631 assert_eq!(rows.len(), 1);
3632 assert_eq!(rows[0].values[1], SqlValue::Double(25.0));
3633 assert_eq!(rows[0].values[2], SqlValue::Text("a|b".into()));
3634 }
3635
3636 #[test]
3637 fn partial_state_matches_single_for_count_sum_total_avg_min_max() {
3638 assert_single_equals_merged(
3639 || Box::new(CountAccumulator::new(false)),
3640 vec![
3641 vec![
3642 Some(SqlValue::Integer(1)),
3643 Some(SqlValue::BigInt(2)),
3644 Some(SqlValue::Text("x".into())),
3645 Some(SqlValue::Null),
3646 ],
3647 vec![Some(SqlValue::Integer(3))],
3648 ],
3649 );
3650 assert_single_equals_merged(
3651 || Box::new(SumAccumulator::new()),
3652 vec![
3653 vec![
3654 Some(SqlValue::Integer(1)),
3655 Some(SqlValue::BigInt(2)),
3656 Some(SqlValue::Float(3.5)),
3657 ],
3658 vec![Some(SqlValue::Double(4.5)), Some(SqlValue::Null)],
3659 ],
3660 );
3661 assert_single_equals_merged(
3662 || Box::new(TotalAccumulator::new()),
3663 vec![
3664 vec![Some(SqlValue::Integer(1)), Some(SqlValue::Null)],
3665 vec![Some(SqlValue::Double(2.5))],
3666 ],
3667 );
3668 assert_single_equals_merged(
3669 || Box::new(AvgAccumulator::new()),
3670 vec![
3671 vec![Some(SqlValue::Integer(2)), Some(SqlValue::Double(4.0))],
3672 vec![Some(SqlValue::Null), Some(SqlValue::Double(6.0))],
3673 ],
3674 );
3675 assert_single_equals_merged(
3676 || Box::new(MinMaxAccumulator::new(true)),
3677 vec![
3678 vec![Some(SqlValue::Integer(3)), Some(SqlValue::Integer(1))],
3679 vec![Some(SqlValue::Integer(2)), Some(SqlValue::Null)],
3680 ],
3681 );
3682 assert_single_equals_merged(
3683 || Box::new(MinMaxAccumulator::new(false)),
3684 vec![
3685 vec![
3686 Some(SqlValue::Text("b".into())),
3687 Some(SqlValue::Text("a".into())),
3688 ],
3689 vec![Some(SqlValue::Text("c".into())), Some(SqlValue::Null)],
3690 ],
3691 );
3692 }
3693
3694 #[test]
3695 fn partial_state_matches_single_for_ordered_string_aggregates() {
3696 assert_single_equals_merged(
3697 || Box::new(GroupConcatAccumulator::new("|".into())),
3698 vec![
3699 vec![Some(SqlValue::Text("a".into())), Some(SqlValue::Null)],
3700 vec![
3701 Some(SqlValue::Text("b".into())),
3702 Some(SqlValue::Text("c".into())),
3703 ],
3704 ],
3705 );
3706 assert_single_equals_merged(
3707 || Box::new(StringAggAccumulator::new("::".into())),
3708 vec![
3709 vec![Some(SqlValue::Text("a".into()))],
3710 vec![Some(SqlValue::Text("b".into())), Some(SqlValue::Null)],
3711 ],
3712 );
3713 }
3714
3715 #[test]
3716 fn partial_state_handles_empty_all_null_single_and_mixed_boundaries() {
3717 assert_single_equals_merged(
3718 || Box::new(CountAccumulator::new(false)),
3719 vec![vec![], vec![]],
3720 );
3721 assert_single_equals_merged(|| Box::new(SumAccumulator::new()), vec![vec![], vec![]]);
3722 assert_single_equals_merged(|| Box::new(TotalAccumulator::new()), vec![vec![], vec![]]);
3723 assert_single_equals_merged(|| Box::new(AvgAccumulator::new()), vec![vec![], vec![]]);
3724 assert_single_equals_merged(
3725 || Box::new(MinMaxAccumulator::new(true)),
3726 vec![vec![], vec![]],
3727 );
3728 assert_single_equals_merged(
3729 || Box::new(MinMaxAccumulator::new(false)),
3730 vec![vec![], vec![]],
3731 );
3732 assert_single_equals_merged(
3733 || Box::new(GroupConcatAccumulator::new(",".into())),
3734 vec![vec![], vec![]],
3735 );
3736 assert_single_equals_merged(
3737 || Box::new(StringAggAccumulator::new(",".into())),
3738 vec![vec![], vec![]],
3739 );
3740 assert_single_equals_merged(
3741 || Box::new(CountAccumulator::new(false)),
3742 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3743 );
3744 assert_single_equals_merged(
3745 || Box::new(SumAccumulator::new()),
3746 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3747 );
3748 assert_single_equals_merged(
3749 || Box::new(TotalAccumulator::new()),
3750 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3751 );
3752 assert_single_equals_merged(
3753 || Box::new(AvgAccumulator::new()),
3754 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3755 );
3756 assert_single_equals_merged(
3757 || Box::new(MinMaxAccumulator::new(true)),
3758 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3759 );
3760 assert_single_equals_merged(
3761 || Box::new(MinMaxAccumulator::new(false)),
3762 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3763 );
3764 assert_single_equals_merged(
3765 || Box::new(GroupConcatAccumulator::new(",".into())),
3766 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3767 );
3768 assert_single_equals_merged(
3769 || Box::new(StringAggAccumulator::new(",".into())),
3770 vec![vec![Some(SqlValue::Null)], vec![Some(SqlValue::Null)]],
3771 );
3772 assert_single_equals_merged(
3773 || Box::new(SumAccumulator::new()),
3774 vec![vec![Some(SqlValue::Integer(7))]],
3775 );
3776 assert_single_equals_merged(
3777 || Box::new(AvgAccumulator::new()),
3778 vec![
3779 vec![Some(SqlValue::Null)],
3780 vec![Some(SqlValue::Double(8.0))],
3781 ],
3782 );
3783 }
3784
3785 #[test]
3786 fn commutative_accumulators_are_merge_order_invariant() {
3787 let orders = vec![
3788 vec![1, 3, 0, 2],
3789 vec![3, 2, 1, 0],
3790 vec![0, 1, 2, 3],
3791 vec![2, 0, 3, 1],
3792 ];
3793 assert_merge_order_invariant(
3794 || Box::new(CountAccumulator::new(false)),
3795 vec![vec![Some(SqlValue::Integer(1))]],
3796 &[vec![0]],
3797 );
3798 assert_merge_order_invariant(
3799 || Box::new(SumAccumulator::new()),
3800 vec![
3801 vec![Some(SqlValue::Integer(1))],
3802 vec![Some(SqlValue::Integer(2))],
3803 ],
3804 &[vec![0, 1], vec![1, 0]],
3805 );
3806 assert_merge_order_invariant(
3807 || Box::new(AvgAccumulator::new()),
3808 vec![
3809 vec![Some(SqlValue::Integer(1))],
3810 vec![Some(SqlValue::Integer(2))],
3811 vec![Some(SqlValue::Integer(3))],
3812 ],
3813 &[vec![0, 1, 2], vec![2, 1, 0]],
3814 );
3815 let numeric_partitions = vec![
3816 vec![Some(SqlValue::Integer(1)), Some(SqlValue::Null)],
3817 vec![Some(SqlValue::BigInt(2))],
3818 vec![],
3819 vec![Some(SqlValue::Double(3.0))],
3820 ];
3821 assert_merge_order_invariant(
3822 || Box::new(CountAccumulator::new(false)),
3823 numeric_partitions.clone(),
3824 &orders,
3825 );
3826 assert_merge_order_invariant(
3827 || Box::new(SumAccumulator::new()),
3828 numeric_partitions.clone(),
3829 &orders,
3830 );
3831 assert_merge_order_invariant(
3832 || Box::new(TotalAccumulator::new()),
3833 numeric_partitions.clone(),
3834 &orders,
3835 );
3836 assert_merge_order_invariant(
3837 || Box::new(AvgAccumulator::new()),
3838 numeric_partitions.clone(),
3839 &orders,
3840 );
3841 let integer_partitions = vec![
3842 vec![Some(SqlValue::Integer(3)), Some(SqlValue::Null)],
3843 vec![Some(SqlValue::Integer(1))],
3844 vec![],
3845 vec![Some(SqlValue::Integer(2))],
3846 ];
3847 assert_merge_order_invariant(
3848 || Box::new(MinMaxAccumulator::new(true)),
3849 integer_partitions.clone(),
3850 &orders,
3851 );
3852 assert_merge_order_invariant(
3853 || Box::new(MinMaxAccumulator::new(false)),
3854 integer_partitions,
3855 &orders,
3856 );
3857 }
3858
3859 #[test]
3860 fn avg_partial_state_uses_sum_count_and_never_divides_by_zero_during_merge() {
3861 let empty = {
3862 let acc = AvgAccumulator::new();
3863 acc.state().unwrap()
3864 };
3865 assert_eq!(empty, vec![SqlValue::Double(0.0), SqlValue::BigInt(0)]);
3866
3867 let mut partial = AvgAccumulator::new();
3868 partial.update(Some(SqlValue::Integer(2))).unwrap();
3869 partial.update(Some(SqlValue::Double(4.0))).unwrap();
3870 assert_eq!(
3871 partial.state().unwrap(),
3872 vec![SqlValue::Double(6.0), SqlValue::BigInt(2)]
3873 );
3874
3875 let mut final_acc = AvgAccumulator::new();
3876 final_acc.merge(&empty).unwrap();
3877 assert_eq!(final_acc.finalize().unwrap(), SqlValue::Null);
3878 final_acc.merge(&partial.state().unwrap()).unwrap();
3879 assert_eq!(final_acc.finalize().unwrap(), SqlValue::Double(3.0));
3880 }
3881
3882 #[test]
3883 fn merge_rejects_invalid_state_contracts_without_panicking() {
3884 let mut count = CountAccumulator::new(false);
3885 assert!(count.merge(&[]).is_err());
3886 assert!(count.merge(&[SqlValue::Text("bad".into())]).is_err());
3887
3888 let mut avg = AvgAccumulator::new();
3889 assert!(avg.merge(&[SqlValue::Double(1.0)]).is_err());
3890 assert!(
3891 avg.merge(&[SqlValue::Double(1.0), SqlValue::Text("bad".into())])
3892 .is_err()
3893 );
3894
3895 let mut concat = GroupConcatAccumulator::new("|".into());
3896 assert!(
3897 concat
3898 .merge(&[SqlValue::Text("a".into()), SqlValue::Text(",".into())])
3899 .is_err()
3900 );
3901 }
3902
3903 #[test]
3904 fn count_accumulator_counts_rows_and_skips_nulls() {
3905 let mut acc = CountAccumulator::new(false);
3906 acc.update(None).unwrap();
3907 acc.update(Some(SqlValue::Null)).unwrap();
3908 acc.update(Some(SqlValue::Integer(1))).unwrap();
3909 assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(2));
3910 }
3911
3912 #[test]
3913 fn count_accumulator_distinct_deduplicates() {
3914 let mut acc = CountAccumulator::new(true);
3915 acc.update(Some(SqlValue::Integer(1))).unwrap();
3916 acc.update(Some(SqlValue::Integer(1))).unwrap();
3917 acc.update(Some(SqlValue::Integer(2))).unwrap();
3918 assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(2));
3919 }
3920
3921 #[test]
3922 fn count_distinct_uses_group_key_equality_boundaries() {
3923 let mut acc = CountAccumulator::new(true);
3924 let nan_a = f64::from_bits(0x7ff8_0000_0000_0001);
3925 let nan_b = f64::from_bits(0x7ff8_0000_0000_0002);
3926 for value in [
3927 SqlValue::Null,
3928 SqlValue::Null,
3929 SqlValue::Integer(1),
3930 SqlValue::Integer(1),
3931 SqlValue::Double(1.0),
3932 SqlValue::Double(-0.0),
3933 SqlValue::Double(0.0),
3934 SqlValue::Double(nan_a),
3935 SqlValue::Double(nan_a),
3936 SqlValue::Double(nan_b),
3937 SqlValue::Text("same".into()),
3938 SqlValue::Text("same".into()),
3939 SqlValue::Blob(vec![1, 2]),
3940 SqlValue::Blob(vec![1, 2]),
3941 SqlValue::Blob(vec![1, 3]),
3942 ] {
3943 acc.update(Some(value)).unwrap();
3944 }
3945 assert_eq!(acc.finalize().unwrap(), SqlValue::BigInt(9));
3946 }
3947
3948 #[test]
3949 fn distinct_non_count_accumulators_deduplicate_non_null_values() {
3950 let mut sum = SumAccumulator::with_distinct(true);
3951 for value in [
3952 SqlValue::Integer(1),
3953 SqlValue::Integer(1),
3954 SqlValue::Double(1.0),
3955 SqlValue::Integer(2),
3956 SqlValue::Null,
3957 ] {
3958 sum.update(Some(value)).unwrap();
3959 }
3960 assert_eq!(sum.finalize().unwrap(), SqlValue::Double(4.0));
3961
3962 let mut avg = AvgAccumulator::with_distinct(true);
3963 for value in [
3964 SqlValue::Integer(1),
3965 SqlValue::Integer(1),
3966 SqlValue::Integer(3),
3967 SqlValue::Null,
3968 ] {
3969 avg.update(Some(value)).unwrap();
3970 }
3971 assert_eq!(avg.finalize().unwrap(), SqlValue::Double(2.0));
3972
3973 let mut min = MinMaxAccumulator::with_distinct(true, true);
3974 let mut max = MinMaxAccumulator::with_distinct(false, true);
3975 for value in [
3976 SqlValue::Text("b".into()),
3977 SqlValue::Text("a".into()),
3978 SqlValue::Text("a".into()),
3979 SqlValue::Text("c".into()),
3980 ] {
3981 min.update(Some(value.clone())).unwrap();
3982 max.update(Some(value)).unwrap();
3983 }
3984 assert_eq!(min.finalize().unwrap(), SqlValue::Text("a".into()));
3985 assert_eq!(max.finalize().unwrap(), SqlValue::Text("c".into()));
3986
3987 let mut group_concat = GroupConcatAccumulator::with_distinct("|".into(), true);
3988 let mut string_agg = StringAggAccumulator::with_distinct(";".into(), true);
3989 for value in [
3990 SqlValue::Text("a".into()),
3991 SqlValue::Text("a".into()),
3992 SqlValue::Null,
3993 SqlValue::Text("b".into()),
3994 ] {
3995 group_concat.update(Some(value.clone())).unwrap();
3996 string_agg.update(Some(value)).unwrap();
3997 }
3998 assert_eq!(
3999 group_concat.finalize().unwrap(),
4000 SqlValue::Text("a|b".into())
4001 );
4002 assert_eq!(string_agg.finalize().unwrap(), SqlValue::Text("a;b".into()));
4003 }
4004
4005 #[test]
4006 fn sum_accumulator_aggregates_numeric_values() {
4007 let mut acc = SumAccumulator::new();
4008 acc.update(Some(SqlValue::Integer(2))).unwrap();
4009 acc.update(Some(SqlValue::Double(3.5))).unwrap();
4010 acc.update(Some(SqlValue::Null)).unwrap();
4011 assert_eq!(acc.finalize().unwrap(), SqlValue::Double(5.5));
4012 }
4013
4014 #[test]
4015 fn total_accumulator_returns_zero_for_empty() {
4016 let acc = TotalAccumulator::new();
4017 assert_eq!(acc.finalize().unwrap(), SqlValue::Double(0.0));
4018 }
4019
4020 #[test]
4021 fn total_accumulator_aggregates_numeric_values() {
4022 let mut acc = TotalAccumulator::new();
4023 acc.update(Some(SqlValue::Integer(2))).unwrap();
4024 acc.update(Some(SqlValue::Null)).unwrap();
4025 acc.update(Some(SqlValue::Double(1.5))).unwrap();
4026 assert_eq!(acc.finalize().unwrap(), SqlValue::Double(3.5));
4027 }
4028
4029 #[test]
4030 fn avg_accumulator_handles_empty_and_nulls() {
4031 let mut acc = AvgAccumulator::new();
4032 assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
4033 acc.update(Some(SqlValue::Null)).unwrap();
4034 acc.update(Some(SqlValue::BigInt(4))).unwrap();
4035 acc.update(Some(SqlValue::Integer(2))).unwrap();
4036 assert_eq!(acc.finalize().unwrap(), SqlValue::Double(3.0));
4037 }
4038
4039 #[test]
4040 fn min_max_accumulator_tracks_extremes() {
4041 let mut min_acc = MinMaxAccumulator::new(true);
4042 let mut max_acc = MinMaxAccumulator::new(false);
4043 for value in [3, 1, 2] {
4044 min_acc.update(Some(SqlValue::Integer(value))).unwrap();
4045 max_acc.update(Some(SqlValue::Integer(value))).unwrap();
4046 }
4047 assert_eq!(min_acc.finalize().unwrap(), SqlValue::Integer(1));
4048 assert_eq!(max_acc.finalize().unwrap(), SqlValue::Integer(3));
4049 }
4050
4051 #[test]
4052 fn min_max_accumulator_rejects_type_mismatch() {
4053 let mut acc = MinMaxAccumulator::new(true);
4054 acc.update(Some(SqlValue::Integer(1))).unwrap();
4055 let err = acc.update(Some(SqlValue::Text("bad".into()))).unwrap_err();
4056 match err {
4057 ExecutorError::Evaluation(crate::executor::EvaluationError::TypeMismatch {
4058 ..
4059 }) => {}
4060 other => panic!("unexpected error {:?}", other),
4061 }
4062 }
4063
4064 #[test]
4065 fn group_concat_accumulator_joins_values() {
4066 let mut acc = GroupConcatAccumulator::new("|".into());
4067 acc.update(Some(SqlValue::Text("a".into()))).unwrap();
4068 acc.update(Some(SqlValue::Null)).unwrap();
4069 acc.update(Some(SqlValue::Text("b".into()))).unwrap();
4070 assert_eq!(acc.finalize().unwrap(), SqlValue::Text("a|b".into()));
4071 }
4072
4073 #[test]
4074 fn group_concat_accumulator_empty_returns_null() {
4075 let acc = GroupConcatAccumulator::new(",".into());
4076 assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
4077 }
4078
4079 #[test]
4080 fn string_agg_accumulator_joins_values() {
4081 let mut acc = StringAggAccumulator::new("::".into());
4082 acc.update(Some(SqlValue::Text("a".into()))).unwrap();
4083 acc.update(Some(SqlValue::Null)).unwrap();
4084 acc.update(Some(SqlValue::Text("b".into()))).unwrap();
4085 assert_eq!(acc.finalize().unwrap(), SqlValue::Text("a::b".into()));
4086 }
4087
4088 #[test]
4089 fn string_agg_accumulator_empty_returns_null() {
4090 let acc = StringAggAccumulator::new(",".into());
4091 assert_eq!(acc.finalize().unwrap(), SqlValue::Null);
4092 }
4093
4094 #[test]
4095 fn encode_group_key_is_deterministic() {
4096 let values = vec![
4097 SqlValue::Integer(1),
4098 SqlValue::Text("a".into()),
4099 SqlValue::Null,
4100 ];
4101 let first = encode_group_key(&values).unwrap();
4102 let second = encode_group_key(&values).unwrap();
4103 assert_eq!(first, second);
4104 }
4105
4106 #[test]
4107 fn percentile_disc_accumulator_follows_postgres_selection_rule() {
4108 let values = [1i64, 2, 2, 3];
4110 for (fraction, expected) in [(0.0, 1i64), (0.25, 1), (0.5, 2), (0.75, 2), (1.0, 3)] {
4111 let mut acc = PercentileDiscAccumulator::new(fraction, vec![(true, false)]);
4112 for value in values {
4113 let value = SqlValue::Integer(value as i32);
4114 acc.update_ordered(Some(value.clone()), std::slice::from_ref(&value))
4115 .unwrap();
4116 }
4117 acc.update_ordered(Some(SqlValue::Null), &[SqlValue::Null])
4119 .unwrap();
4120 assert_eq!(
4121 acc.finalize().unwrap(),
4122 SqlValue::Integer(expected as i32),
4123 "fraction {fraction}"
4124 );
4125 }
4126
4127 let empty = PercentileDiscAccumulator::new(0.5, vec![(true, false)]);
4128 assert_eq!(empty.finalize().unwrap(), SqlValue::Null);
4129
4130 let acc = PercentileDiscAccumulator::new(0.5, vec![(true, false)]);
4131 assert!(acc.state().is_err(), "ordered-set partial state is invalid");
4132 }
4133
4134 #[test]
4135 fn ordered_string_accumulators_sort_stably_and_reject_partial_state() {
4136 let mut acc = StringAggAccumulator::with_order(",".into(), false, vec![(false, false)]);
4137 for (key, value) in [(2, "b"), (1, "a"), (3, "c"), (2, "d")] {
4138 acc.update_ordered(
4139 Some(SqlValue::Text(value.into())),
4140 &[SqlValue::Integer(key)],
4141 )
4142 .unwrap();
4143 }
4144 assert_eq!(acc.finalize().unwrap(), SqlValue::Text("c,b,d,a".into()));
4146 assert!(acc.state().is_err());
4147 assert!(acc.merge(&[]).is_err());
4148
4149 let mut concat = GroupConcatAccumulator::with_order("|".into(), true, vec![(true, false)]);
4150 for (key, value) in [(2, "x"), (1, "y"), (3, "x")] {
4151 concat
4152 .update_ordered(
4153 Some(SqlValue::Text(value.into())),
4154 &[SqlValue::Integer(key)],
4155 )
4156 .unwrap();
4157 }
4158 assert_eq!(concat.finalize().unwrap(), SqlValue::Text("y|x".into()));
4160 }
4161
4162 #[test]
4163 fn stable_moments_merge_and_preserve_sample_boundaries() {
4164 let mut whole = StatisticsAccumulator::new(StatisticsKind::Variance(true));
4165 let mut left = whole.clone();
4166 let mut right = whole.clone();
4167 for (index, value) in [1.0e12, 1.0e12 + 1.0, 1.0e12 + 2.0, 1.0e12 + 3.0]
4168 .into_iter()
4169 .enumerate()
4170 {
4171 whole.update(Some(SqlValue::Double(value))).unwrap();
4172 if index < 2 {
4173 left.update(Some(SqlValue::Double(value))).unwrap();
4174 } else {
4175 right.update(Some(SqlValue::Double(value))).unwrap();
4176 }
4177 }
4178 left.merge(&right.state().unwrap()).unwrap();
4179 assert_eq!(whole.finalize().unwrap(), left.finalize().unwrap());
4180 assert_eq!(whole.finalize().unwrap(), SqlValue::Double(5.0 / 3.0));
4181
4182 let empty = StatisticsAccumulator::new(StatisticsKind::Variance(true));
4183 assert_eq!(empty.finalize().unwrap(), SqlValue::Null);
4184 let mut singleton = empty.clone();
4185 singleton.update(Some(SqlValue::Integer(7))).unwrap();
4186 assert_eq!(singleton.finalize().unwrap(), SqlValue::Null);
4187 singleton.kind = StatisticsKind::Variance(false);
4188 assert_eq!(singleton.finalize().unwrap(), SqlValue::Double(0.0));
4189 }
4190
4191 #[test]
4192 fn continuous_percentiles_preserve_non_finite_endpoints() {
4193 for (fraction, expected) in [(0.0, f64::NEG_INFINITY), (1.0, f64::INFINITY)] {
4194 let mut acc = PercentileContAccumulator::new(fraction, true);
4195 acc.update(Some(SqlValue::Double(f64::NEG_INFINITY)))
4196 .unwrap();
4197 acc.update(Some(SqlValue::Double(0.0))).unwrap();
4198 acc.update(Some(SqlValue::Double(f64::INFINITY))).unwrap();
4199 assert_eq!(acc.finalize().unwrap(), SqlValue::Double(expected));
4200 }
4201
4202 let mut acc = PercentileContAccumulator::new(1.0, true);
4203 acc.update(Some(SqlValue::Double(f64::NAN))).unwrap();
4204 let SqlValue::Double(value) = acc.finalize().unwrap() else {
4205 panic!("continuous percentile must return DOUBLE");
4206 };
4207 assert!(value.is_nan());
4208 }
4209}