1use arrow::array::{
21 Array, ArrayRef, ArrowNativeTypeOp, ArrowNumericType, ArrowPrimitiveType, AsArray,
22 BooleanArray, PrimitiveArray, PrimitiveBuilder, UInt64Array,
23};
24
25use arrow::compute::{DecimalCast, sum};
26use arrow::datatypes::{
27 ArrowNativeType, DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE,
28 DECIMAL64_MAX_PRECISION, DECIMAL64_MAX_SCALE, DECIMAL128_MAX_PRECISION,
29 DECIMAL128_MAX_SCALE, DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, DataType,
30 Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType,
31 DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
32 DurationSecondType, Field, FieldRef, Float64Type, TimeUnit, UInt64Type,
33};
34use datafusion_common::types::{NativeType, logical_float64};
35use datafusion_common::{
36 Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, not_impl_err,
37};
38use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
39use datafusion_expr::utils::format_state_name;
40use datafusion_expr::{
41 Accumulator, AggregateUDFImpl, Coercion, Documentation, EmitTo, Expr,
42 GroupsAccumulator, ReversedUDAF, Signature, TypeSignature, TypeSignatureClass,
43 Volatility,
44};
45use datafusion_functions_aggregate_common::aggregate::avg_distinct::{
46 DecimalDistinctAvgAccumulator, Float64DistinctAvgAccumulator,
47};
48use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::NullState;
49use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{
50 filtered_null_mask, set_nulls,
51};
52use datafusion_functions_aggregate_common::utils::DecimalAverager;
53use datafusion_macros::user_doc;
54use log::debug;
55use std::fmt::Debug;
56use std::marker::PhantomData;
57use std::mem::{size_of, size_of_val};
58use std::sync::Arc;
59
60make_udaf_expr_and_func!(
61 Avg,
62 avg,
63 expression,
64 "Returns the avg of a group of values.",
65 avg_udaf
66);
67
68pub fn avg_distinct(expr: Expr) -> Expr {
69 Expr::AggregateFunction(datafusion_expr::expr::AggregateFunction::new_udf(
70 avg_udaf(),
71 vec![expr],
72 true,
73 None,
74 vec![],
75 None,
76 ))
77}
78
79#[user_doc(
80 doc_section(label = "General Functions"),
81 description = "Returns the average of numeric values in the specified column.",
82 syntax_example = "avg(expression)",
83 sql_example = r#"```sql
84> SELECT avg(column_name) FROM table_name;
85+---------------------------+
86| avg(column_name) |
87+---------------------------+
88| 42.75 |
89+---------------------------+
90```"#,
91 standard_argument(name = "expression",)
92)]
93#[derive(Debug, PartialEq, Eq, Hash)]
94pub struct Avg {
95 signature: Signature,
96 aliases: Vec<String>,
97}
98
99impl Avg {
100 pub fn new() -> Self {
101 Self {
102 signature: Signature::one_of(
105 vec![
106 TypeSignature::Coercible(vec![Coercion::new_exact(
107 TypeSignatureClass::Decimal,
108 )]),
109 TypeSignature::Coercible(vec![Coercion::new_exact(
110 TypeSignatureClass::Duration,
111 )]),
112 TypeSignature::Coercible(vec![Coercion::new_implicit(
113 TypeSignatureClass::Native(logical_float64()),
114 vec![TypeSignatureClass::Integer, TypeSignatureClass::Float],
115 NativeType::Float64,
116 )]),
117 ],
118 Volatility::Immutable,
119 ),
120 aliases: vec![String::from("mean")],
121 }
122 }
123}
124
125impl Default for Avg {
126 fn default() -> Self {
127 Self::new()
128 }
129}
130
131const AVG_SUM_HEADROOM_DIGITS: u8 = 13;
142
143fn avg_sum_data_type(data_type: &DataType) -> DataType {
146 let (precision, scale, input_max_precision) = match data_type {
147 DataType::Decimal32(precision, scale) => {
148 (*precision, *scale, DECIMAL32_MAX_PRECISION)
149 }
150 DataType::Decimal64(precision, scale) => {
151 (*precision, *scale, DECIMAL64_MAX_PRECISION)
152 }
153 DataType::Decimal128(precision, scale) => {
154 (*precision, *scale, DECIMAL128_MAX_PRECISION)
155 }
156 DataType::Decimal256(precision, scale) => {
157 (*precision, *scale, DECIMAL256_MAX_PRECISION)
158 }
159 data_type => return data_type.clone(),
160 };
161
162 let required = precision
163 .saturating_add(AVG_SUM_HEADROOM_DIGITS)
164 .max(input_max_precision);
165
166 if required <= DECIMAL64_MAX_PRECISION {
169 DataType::Decimal64(DECIMAL64_MAX_PRECISION, scale)
170 } else if required <= DECIMAL128_MAX_PRECISION {
171 DataType::Decimal128(DECIMAL128_MAX_PRECISION, scale)
172 } else {
173 DataType::Decimal256(DECIMAL256_MAX_PRECISION, scale)
174 }
175}
176
177macro_rules! decimal_avg_dispatch {
180 ($input:expr, $sum:expr, $builder:ident, $($arg:expr),*) => {
181 match ($input, $sum) {
182 (DataType::Decimal32(..), DataType::Decimal64(..)) => {
183 $builder::<Decimal32Type, Decimal64Type>($($arg),*)
184 }
185 (DataType::Decimal32(..), DataType::Decimal128(..)) => {
186 $builder::<Decimal32Type, Decimal128Type>($($arg),*)
187 }
188 (DataType::Decimal64(..), DataType::Decimal64(..)) => {
189 $builder::<Decimal64Type, Decimal64Type>($($arg),*)
190 }
191 (DataType::Decimal64(..), DataType::Decimal128(..)) => {
192 $builder::<Decimal64Type, Decimal128Type>($($arg),*)
193 }
194 (DataType::Decimal128(..), DataType::Decimal128(..)) => {
195 $builder::<Decimal128Type, Decimal128Type>($($arg),*)
196 }
197 (DataType::Decimal128(..), DataType::Decimal256(..)) => {
198 $builder::<Decimal128Type, Decimal256Type>($($arg),*)
199 }
200 (DataType::Decimal256(..), DataType::Decimal256(..)) => {
201 $builder::<Decimal256Type, Decimal256Type>($($arg),*)
202 }
203 (input, sum) => {
204 internal_err!("avg cannot accumulate {input} as {sum}")
205 }
206 }
207 };
208}
209
210impl AggregateUDFImpl for Avg {
211 fn name(&self) -> &str {
212 "avg"
213 }
214
215 fn signature(&self) -> &Signature {
216 &self.signature
217 }
218
219 fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
220 match &arg_types[0] {
221 DataType::Decimal32(precision, scale) => {
222 let new_precision = DECIMAL32_MAX_PRECISION.min(*precision + 4);
225 let new_scale = DECIMAL32_MAX_SCALE.min(*scale + 4);
226 Ok(DataType::Decimal32(new_precision, new_scale))
227 }
228 DataType::Decimal64(precision, scale) => {
229 let new_precision = DECIMAL64_MAX_PRECISION.min(*precision + 4);
232 let new_scale = DECIMAL64_MAX_SCALE.min(*scale + 4);
233 Ok(DataType::Decimal64(new_precision, new_scale))
234 }
235 DataType::Decimal128(precision, scale) => {
236 let new_precision = DECIMAL128_MAX_PRECISION.min(*precision + 4);
239 let new_scale = DECIMAL128_MAX_SCALE.min(*scale + 4);
240 Ok(DataType::Decimal128(new_precision, new_scale))
241 }
242 DataType::Decimal256(precision, scale) => {
243 let new_precision = DECIMAL256_MAX_PRECISION.min(*precision + 4);
246 let new_scale = DECIMAL256_MAX_SCALE.min(*scale + 4);
247 Ok(DataType::Decimal256(new_precision, new_scale))
248 }
249 DataType::Duration(time_unit) => Ok(DataType::Duration(*time_unit)),
250 _ => Ok(DataType::Float64),
251 }
252 }
253
254 fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
255 let data_type = acc_args.expr_fields[0].data_type();
256 use DataType::*;
257
258 if acc_args.is_distinct {
260 match (data_type, acc_args.return_type()) {
261 (Float64, _) => Ok(Box::new(Float64DistinctAvgAccumulator::default())),
263
264 (Decimal32(..), Decimal32(..))
265 | (Decimal64(..), Decimal64(..))
266 | (Decimal128(..), Decimal128(..))
267 | (Decimal256(..), Decimal256(..)) => {
268 let sum_data_type = avg_sum_data_type(data_type);
269 decimal_avg_dispatch!(
270 data_type,
271 &sum_data_type,
272 decimal_distinct_avg_accumulator,
273 &sum_data_type,
274 acc_args.return_type()
275 )
276 }
277
278 (dt, return_type) => exec_err!(
279 "AVG(DISTINCT) for ({} --> {}) not supported",
280 dt,
281 return_type
282 ),
283 }
284 } else {
285 match (&data_type, acc_args.return_type()) {
286 (Float64, Float64) => Ok(Box::<AvgAccumulator>::default()),
287 (Decimal32(..), Decimal32(..))
288 | (Decimal64(..), Decimal64(..))
289 | (Decimal128(..), Decimal128(..))
290 | (Decimal256(..), Decimal256(..)) => {
291 let sum_data_type = avg_sum_data_type(data_type);
292 decimal_avg_dispatch!(
293 data_type,
294 &sum_data_type,
295 decimal_avg_accumulator,
296 sum_data_type.clone(),
297 acc_args.return_type().clone()
298 )
299 }
300
301 (Duration(time_unit), Duration(result_unit)) => {
302 Ok(Box::new(DurationAvgAccumulator {
303 sum: None,
304 count: 0,
305 time_unit: *time_unit,
306 result_unit: *result_unit,
307 }))
308 }
309
310 (dt, return_type) => {
311 exec_err!("AvgAccumulator for ({} --> {})", dt, return_type)
312 }
313 }
314 }
315 }
316
317 fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
318 if args.is_distinct {
319 let dt = match args.input_fields[0].data_type() {
322 DataType::Decimal32(_, scale) => {
323 DataType::Decimal32(DECIMAL32_MAX_PRECISION, *scale)
324 }
325 DataType::Decimal64(_, scale) => {
326 DataType::Decimal64(DECIMAL64_MAX_PRECISION, *scale)
327 }
328 DataType::Decimal128(_, scale) => {
329 DataType::Decimal128(DECIMAL128_MAX_PRECISION, *scale)
330 }
331 DataType::Decimal256(_, scale) => {
332 DataType::Decimal256(DECIMAL256_MAX_PRECISION, *scale)
333 }
334 _ => args.return_type().clone(),
335 };
336 Ok(vec![
339 Field::new_list(
340 format_state_name(args.name, "avg distinct"),
341 Field::new_list_field(dt, true),
342 false,
343 )
344 .into(),
345 ])
346 } else {
347 let sum_data_type = avg_sum_data_type(args.input_fields[0].data_type());
348 Ok(vec![
349 Field::new(
350 format_state_name(args.name, "count"),
351 DataType::UInt64,
352 true,
353 ),
354 Field::new(format_state_name(args.name, "sum"), sum_data_type, true),
355 ]
356 .into_iter()
357 .map(Arc::new)
358 .collect())
359 }
360 }
361
362 fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
363 matches!(
364 args.return_field.data_type(),
365 DataType::Float64
366 | DataType::Decimal32(_, _)
367 | DataType::Decimal64(_, _)
368 | DataType::Decimal128(_, _)
369 | DataType::Decimal256(_, _)
370 | DataType::Duration(_)
371 ) && !args.is_distinct
372 }
373
374 fn create_groups_accumulator(
375 &self,
376 args: AccumulatorArgs,
377 ) -> Result<Box<dyn GroupsAccumulator>> {
378 use DataType::*;
379
380 let data_type = args.expr_fields[0].data_type();
381
382 match (data_type, args.return_field.data_type()) {
384 (Float64, Float64) => {
385 Ok(Box::new(AvgGroupsAccumulator::<Float64Type, _>::new(
386 data_type,
387 args.return_field.data_type(),
388 |sum: f64, count: u64| Ok(sum / count as f64),
389 )))
390 }
391 (Decimal32(..), Decimal32(..))
392 | (Decimal64(..), Decimal64(..))
393 | (Decimal128(..), Decimal128(..))
394 | (Decimal256(..), Decimal256(..)) => {
395 let sum_data_type = avg_sum_data_type(data_type);
396 decimal_avg_dispatch!(
397 data_type,
398 &sum_data_type,
399 decimal_avg_groups_accumulator,
400 &sum_data_type,
401 args.return_field.data_type()
402 )
403 }
404
405 (Duration(time_unit), Duration(_result_unit)) => {
406 let avg_fn = move |sum: i64, count: u64| Ok(sum / count as i64);
407
408 match time_unit {
409 TimeUnit::Second => Ok(Box::new(AvgGroupsAccumulator::<
410 DurationSecondType,
411 _,
412 >::new(
413 data_type,
414 args.return_type(),
415 avg_fn,
416 ))),
417 TimeUnit::Millisecond => Ok(Box::new(AvgGroupsAccumulator::<
418 DurationMillisecondType,
419 _,
420 >::new(
421 data_type,
422 args.return_type(),
423 avg_fn,
424 ))),
425 TimeUnit::Microsecond => Ok(Box::new(AvgGroupsAccumulator::<
426 DurationMicrosecondType,
427 _,
428 >::new(
429 data_type,
430 args.return_type(),
431 avg_fn,
432 ))),
433 TimeUnit::Nanosecond => Ok(Box::new(AvgGroupsAccumulator::<
434 DurationNanosecondType,
435 _,
436 >::new(
437 data_type,
438 args.return_type(),
439 avg_fn,
440 ))),
441 }
442 }
443
444 _ => not_impl_err!(
445 "AvgGroupsAccumulator for ({} --> {})",
446 &data_type,
447 args.return_field.data_type()
448 ),
449 }
450 }
451
452 fn aliases(&self) -> &[String] {
453 &self.aliases
454 }
455
456 fn reverse_expr(&self) -> ReversedUDAF {
457 ReversedUDAF::Identical
458 }
459
460 fn documentation(&self) -> Option<&Documentation> {
461 self.doc()
462 }
463}
464
465fn decimal_parts(data_type: &DataType) -> Result<(u8, i8)> {
467 match data_type {
468 DataType::Decimal32(precision, scale)
469 | DataType::Decimal64(precision, scale)
470 | DataType::Decimal128(precision, scale)
471 | DataType::Decimal256(precision, scale) => Ok((*precision, *scale)),
472 data_type => internal_err!("expected a decimal type, got {data_type}"),
473 }
474}
475
476fn decimal_avg_fn<I, S>(
477 sum_scale: i8,
478 target_precision: u8,
479 target_scale: i8,
480) -> Result<impl Fn(S::Native, u64) -> Result<I::Native> + Send + Sync + 'static>
481where
482 I: DecimalType,
483 S: DecimalType,
484 I::Native: DecimalCast,
485 S::Native: DecimalCast,
486{
487 let decimal_averager =
488 DecimalAverager::<S>::try_new(sum_scale, target_precision, target_scale)?;
489
490 Ok(move |sum, count: u64| {
491 let Some(count) = usize::try_from(count).ok().and_then(S::Native::from_usize)
492 else {
493 return exec_err!(
494 "Arithmetic overflow in avg: the row count {count} cannot be \
495 represented in the sum type"
496 );
497 };
498
499 I::Native::from_decimal(decimal_averager.avg(sum, count)?).ok_or_else(|| {
504 exec_datafusion_err!(
505 "Arithmetic overflow in avg: the computed average does not fit \
506 the output type"
507 )
508 })
509 })
510}
511
512fn decimal_avg_accumulator<I, S>(
513 sum_data_type: DataType,
514 return_data_type: DataType,
515) -> Result<Box<dyn Accumulator>>
516where
517 I: DecimalType + ArrowNumericType + Debug + Send + Sync,
518 S: DecimalType + ArrowNumericType + Debug + Send + Sync,
519 I::Native: Into<S::Native> + DecimalCast,
520 S::Native: DecimalCast,
521{
522 let (_, sum_scale) = decimal_parts(&sum_data_type)?;
523 let (target_precision, target_scale) = decimal_parts(&return_data_type)?;
524 let avg_fn = decimal_avg_fn::<I, S>(sum_scale, target_precision, target_scale)?;
525
526 Ok(Box::new(DecimalAvgAccumulator::<I, S, _>::new(
527 sum_data_type,
528 return_data_type,
529 avg_fn,
530 )))
531}
532
533fn decimal_distinct_avg_accumulator<I, S>(
534 sum_data_type: &DataType,
535 return_data_type: &DataType,
536) -> Result<Box<dyn Accumulator>>
537where
538 I: DecimalType + ArrowNumericType + Debug + Send + Sync,
539 S: DecimalType + ArrowNumericType + Debug + Send + Sync,
540 I::Native: Into<S::Native> + DecimalCast,
541 S::Native: DecimalCast,
542{
543 let (_, sum_scale) = decimal_parts(sum_data_type)?;
544 let (target_precision, target_scale) = decimal_parts(return_data_type)?;
545
546 Ok(Box::new(
547 DecimalDistinctAvgAccumulator::<I, S>::with_decimal_params(
548 sum_scale,
549 target_precision,
550 target_scale,
551 ),
552 ))
553}
554
555fn decimal_avg_groups_accumulator<I, S>(
556 sum_data_type: &DataType,
557 return_data_type: &DataType,
558) -> Result<Box<dyn GroupsAccumulator>>
559where
560 I: DecimalType + ArrowNumericType + Debug + Send + Sync,
561 S: DecimalType + ArrowNumericType + Debug + Send + Sync,
562 I::Native: Into<S::Native> + DecimalCast,
563 S::Native: DecimalCast,
564{
565 let (_, sum_scale) = decimal_parts(sum_data_type)?;
566 let (target_precision, target_scale) = decimal_parts(return_data_type)?;
567 let avg_fn = decimal_avg_fn::<I, S>(sum_scale, target_precision, target_scale)?;
568
569 Ok(Box::new(AvgGroupsAccumulator::<I, _, S>::new(
570 sum_data_type,
571 return_data_type,
572 avg_fn,
573 )))
574}
575
576#[derive(Debug, Default)]
578pub struct AvgAccumulator {
579 sum: Option<f64>,
580 count: u64,
581}
582
583impl Accumulator for AvgAccumulator {
584 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
585 let values = values[0].as_primitive::<Float64Type>();
586 self.count += (values.len() - values.null_count()) as u64;
587 if let Some(x) = sum(values) {
588 let v = self.sum.get_or_insert(0.);
589 *v += x;
590 }
591 Ok(())
592 }
593
594 fn evaluate(&mut self) -> Result<ScalarValue> {
595 let avg = if self.count == 0 {
600 None
601 } else {
602 self.sum.map(|f| f / self.count as f64)
603 };
604 Ok(ScalarValue::Float64(avg))
605 }
606
607 fn size(&self) -> usize {
608 size_of_val(self)
609 }
610
611 fn state(&mut self) -> Result<Vec<ScalarValue>> {
612 Ok(vec![
613 ScalarValue::from(self.count),
614 ScalarValue::Float64(self.sum),
615 ])
616 }
617
618 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
619 self.count += sum(states[0].as_primitive::<UInt64Type>()).unwrap_or_default();
621
622 if let Some(x) = sum(states[1].as_primitive::<Float64Type>()) {
624 let v = self.sum.get_or_insert(0.);
625 *v += x;
626 }
627 Ok(())
628 }
629 fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
630 let values = values[0].as_primitive::<Float64Type>();
631 self.count -= (values.len() - values.null_count()) as u64;
632 if let Some(x) = sum(values) {
633 self.sum = Some(self.sum.unwrap() - x);
634 }
635 Ok(())
636 }
637
638 fn supports_retract_batch(&self) -> bool {
639 true
640 }
641}
642
643struct DecimalAvgAccumulator<I, S, F>
648where
649 I: DecimalType + ArrowNumericType + Debug,
650 S: DecimalType + ArrowNumericType + Debug,
651 I::Native: Into<S::Native>,
652 F: Fn(S::Native, u64) -> Result<I::Native>,
653{
654 sum: Option<S::Native>,
655 count: u64,
656 sum_data_type: DataType,
657 return_data_type: DataType,
658 avg_fn: F,
659 _phantom: PhantomData<I>,
660}
661
662impl<I, S, F> Debug for DecimalAvgAccumulator<I, S, F>
663where
664 I: DecimalType + ArrowNumericType + Debug,
665 S: DecimalType + ArrowNumericType + Debug,
666 I::Native: Into<S::Native>,
667 F: Fn(S::Native, u64) -> Result<I::Native>,
668{
669 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
670 f.debug_struct("DecimalAvgAccumulator")
671 .field("sum", &self.sum)
672 .field("count", &self.count)
673 .field("sum_data_type", &self.sum_data_type)
674 .field("return_data_type", &self.return_data_type)
675 .finish_non_exhaustive()
676 }
677}
678
679impl<I, S, F> DecimalAvgAccumulator<I, S, F>
680where
681 I: DecimalType + ArrowNumericType + Debug,
682 S: DecimalType + ArrowNumericType + Debug,
683 I::Native: Into<S::Native>,
684 F: Fn(S::Native, u64) -> Result<I::Native>,
685{
686 fn new(sum_data_type: DataType, return_data_type: DataType, avg_fn: F) -> Self {
687 Self {
688 sum: None,
689 count: 0,
690 sum_data_type,
691 return_data_type,
692 avg_fn,
693 _phantom: PhantomData,
694 }
695 }
696}
697
698fn decimal_sum_as<I, S>(values: &PrimitiveArray<I>) -> Option<S::Native>
704where
705 I: DecimalType + ArrowNumericType,
706 S: DecimalType + ArrowNumericType,
707 I::Native: Into<S::Native>,
708{
709 if values.null_count() == values.len() {
711 return None;
712 }
713
714 let mut sum = S::Native::default();
715 if values.null_count() == 0 {
716 for value in values.values() {
717 sum = sum.add_wrapping((*value).into());
718 }
719 } else {
720 for value in values.iter().flatten() {
721 sum = sum.add_wrapping(value.into());
722 }
723 }
724
725 Some(sum)
726}
727
728impl<I, S, F> Accumulator for DecimalAvgAccumulator<I, S, F>
729where
730 I: DecimalType + ArrowNumericType + Debug,
731 S: DecimalType + ArrowNumericType + Debug,
732 I::Native: Into<S::Native>,
733 F: Fn(S::Native, u64) -> Result<I::Native> + Send + Sync + 'static,
734{
735 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
736 let values = values[0].as_primitive::<I>();
737 self.count += (values.len() - values.null_count()) as u64;
738
739 if let Some(x) = decimal_sum_as::<I, S>(values) {
740 let v = self.sum.unwrap_or_default();
741 self.sum = Some(v.add_wrapping(x));
742 }
743 Ok(())
744 }
745
746 fn evaluate(&mut self) -> Result<ScalarValue> {
747 let v = if self.count == 0 {
751 None
752 } else {
753 self.sum.map(|v| (self.avg_fn)(v, self.count)).transpose()?
754 };
755
756 ScalarValue::new_primitive::<I>(v, &self.return_data_type)
757 }
758
759 fn size(&self) -> usize {
760 size_of_val(self)
761 }
762
763 fn state(&mut self) -> Result<Vec<ScalarValue>> {
764 Ok(vec![
765 ScalarValue::from(self.count),
766 ScalarValue::new_primitive::<S>(self.sum, &self.sum_data_type)?,
767 ])
768 }
769
770 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
771 self.count += sum(states[0].as_primitive::<UInt64Type>()).unwrap_or_default();
773
774 if let Some(x) = sum(states[1].as_primitive::<S>()) {
776 let v = self.sum.unwrap_or_default();
777 self.sum = Some(v.add_wrapping(x));
778 }
779 Ok(())
780 }
781 fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
782 let values = values[0].as_primitive::<I>();
783 self.count -= (values.len() - values.null_count()) as u64;
784 if let Some(x) = decimal_sum_as::<I, S>(values) {
785 let v = self.sum.unwrap_or_default();
786 self.sum = Some(v.sub_wrapping(x));
787 }
788 Ok(())
789 }
790
791 fn supports_retract_batch(&self) -> bool {
792 true
793 }
794}
795
796#[derive(Debug)]
798struct DurationAvgAccumulator {
799 sum: Option<i64>,
800 count: u64,
801 time_unit: TimeUnit,
802 result_unit: TimeUnit,
803}
804
805impl Accumulator for DurationAvgAccumulator {
806 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
807 let array = &values[0];
808 self.count += (array.len() - array.null_count()) as u64;
809
810 let sum_value = match self.time_unit {
811 TimeUnit::Second => sum(array.as_primitive::<DurationSecondType>()),
812 TimeUnit::Millisecond => sum(array.as_primitive::<DurationMillisecondType>()),
813 TimeUnit::Microsecond => sum(array.as_primitive::<DurationMicrosecondType>()),
814 TimeUnit::Nanosecond => sum(array.as_primitive::<DurationNanosecondType>()),
815 };
816
817 if let Some(x) = sum_value {
818 let v = self.sum.get_or_insert(0);
819 *v += x;
820 }
821 Ok(())
822 }
823
824 fn evaluate(&mut self) -> Result<ScalarValue> {
825 let avg = if self.count == 0 {
829 None
830 } else {
831 self.sum.map(|sum| sum / self.count as i64)
832 };
833
834 match self.result_unit {
835 TimeUnit::Second => Ok(ScalarValue::DurationSecond(avg)),
836 TimeUnit::Millisecond => Ok(ScalarValue::DurationMillisecond(avg)),
837 TimeUnit::Microsecond => Ok(ScalarValue::DurationMicrosecond(avg)),
838 TimeUnit::Nanosecond => Ok(ScalarValue::DurationNanosecond(avg)),
839 }
840 }
841
842 fn size(&self) -> usize {
843 size_of_val(self)
844 }
845
846 fn state(&mut self) -> Result<Vec<ScalarValue>> {
847 let duration_value = match self.time_unit {
848 TimeUnit::Second => ScalarValue::DurationSecond(self.sum),
849 TimeUnit::Millisecond => ScalarValue::DurationMillisecond(self.sum),
850 TimeUnit::Microsecond => ScalarValue::DurationMicrosecond(self.sum),
851 TimeUnit::Nanosecond => ScalarValue::DurationNanosecond(self.sum),
852 };
853
854 Ok(vec![ScalarValue::from(self.count), duration_value])
855 }
856
857 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
858 self.count += sum(states[0].as_primitive::<UInt64Type>()).unwrap_or_default();
859
860 let sum_value = match self.time_unit {
861 TimeUnit::Second => sum(states[1].as_primitive::<DurationSecondType>()),
862 TimeUnit::Millisecond => {
863 sum(states[1].as_primitive::<DurationMillisecondType>())
864 }
865 TimeUnit::Microsecond => {
866 sum(states[1].as_primitive::<DurationMicrosecondType>())
867 }
868 TimeUnit::Nanosecond => {
869 sum(states[1].as_primitive::<DurationNanosecondType>())
870 }
871 };
872
873 if let Some(x) = sum_value {
874 let v = self.sum.get_or_insert(0);
875 *v += x;
876 }
877 Ok(())
878 }
879
880 fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
881 let array = &values[0];
882 self.count -= (array.len() - array.null_count()) as u64;
883
884 let sum_value = match self.time_unit {
885 TimeUnit::Second => sum(array.as_primitive::<DurationSecondType>()),
886 TimeUnit::Millisecond => sum(array.as_primitive::<DurationMillisecondType>()),
887 TimeUnit::Microsecond => sum(array.as_primitive::<DurationMicrosecondType>()),
888 TimeUnit::Nanosecond => sum(array.as_primitive::<DurationNanosecondType>()),
889 };
890
891 if let Some(x) = sum_value {
892 self.sum = Some(self.sum.unwrap() - x);
893 }
894 Ok(())
895 }
896
897 fn supports_retract_batch(&self) -> bool {
898 true
899 }
900}
901
902#[derive(Debug)]
911struct AvgGroupsAccumulator<I, F, S = I>
912where
913 I: ArrowNumericType + Send,
914 S: ArrowNumericType + Send,
915 I::Native: Into<S::Native>,
916 F: Fn(S::Native, u64) -> Result<I::Native> + Send + 'static,
917{
918 sum_data_type: DataType,
920
921 return_data_type: DataType,
923
924 counts: Vec<u64>,
926
927 sums: Vec<S::Native>,
929
930 null_state: NullState,
932
933 avg_fn: F,
935
936 _phantom: PhantomData<I>,
937}
938
939impl<I, F, S> AvgGroupsAccumulator<I, F, S>
940where
941 I: ArrowNumericType + Send,
942 S: ArrowNumericType + Send,
943 I::Native: Into<S::Native>,
944 F: Fn(S::Native, u64) -> Result<I::Native> + Send + 'static,
945{
946 pub fn new(sum_data_type: &DataType, return_data_type: &DataType, avg_fn: F) -> Self {
947 debug!(
948 "AvgGroupsAccumulator ({}, sum type: {sum_data_type}) --> {return_data_type}",
949 std::any::type_name::<I>()
950 );
951
952 Self {
953 return_data_type: return_data_type.clone(),
954 sum_data_type: sum_data_type.clone(),
955 counts: vec![],
956 sums: vec![],
957 null_state: NullState::new(),
958 avg_fn,
959 _phantom: PhantomData,
960 }
961 }
962}
963
964impl<I, F, S> GroupsAccumulator for AvgGroupsAccumulator<I, F, S>
965where
966 I: ArrowNumericType + Send,
967 S: ArrowNumericType + Send,
968 I::Native: Into<S::Native>,
969 F: Fn(S::Native, u64) -> Result<I::Native> + Send + 'static,
970{
971 fn update_batch(
972 &mut self,
973 values: &[ArrayRef],
974 group_indices: &[usize],
975 opt_filter: Option<&BooleanArray>,
976 total_num_groups: usize,
977 ) -> Result<()> {
978 assert_eq!(values.len(), 1, "single argument to update_batch");
979 let values = values[0].as_primitive::<I>();
980
981 self.counts.resize(total_num_groups, 0);
983 self.sums.resize(total_num_groups, S::default_value());
984
985 self.null_state.accumulate(
986 group_indices,
987 values,
988 opt_filter,
989 total_num_groups,
990 |group_index, new_value| {
991 let sum = unsafe { self.sums.get_unchecked_mut(group_index) };
993 *sum = sum.add_wrapping(new_value.into());
994
995 self.counts[group_index] += 1;
996 },
997 );
998
999 Ok(())
1000 }
1001
1002 fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
1003 let counts = emit_to.take_needed(&mut self.counts);
1004 let sums = emit_to.take_needed(&mut self.sums);
1005 let nulls = self.null_state.build(emit_to);
1006
1007 if let Some(nulls) = &nulls {
1008 assert_eq!(nulls.len(), sums.len());
1009 }
1010 assert_eq!(counts.len(), sums.len());
1011
1012 let array: PrimitiveArray<I> = if let Some(nulls) = &nulls
1015 && nulls.null_count() > 0
1016 {
1017 let mut builder = PrimitiveBuilder::<I>::with_capacity(nulls.len())
1018 .with_data_type(self.return_data_type.clone());
1019 let iter = sums.into_iter().zip(counts).zip(nulls.iter());
1020
1021 for ((sum, count), is_valid) in iter {
1022 if is_valid {
1023 builder.append_value((self.avg_fn)(sum, count)?)
1024 } else {
1025 builder.append_null();
1026 }
1027 }
1028 builder.finish()
1029 } else {
1030 let averages: Vec<I::Native> = sums
1031 .into_iter()
1032 .zip(counts)
1033 .map(|(sum, count)| (self.avg_fn)(sum, count))
1034 .collect::<Result<Vec<_>>>()?;
1035 PrimitiveArray::new(averages.into(), nulls) .with_data_type(self.return_data_type.clone())
1037 };
1038
1039 Ok(Arc::new(array))
1040 }
1041
1042 fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
1044 let nulls = self.null_state.build(emit_to);
1045
1046 let counts = emit_to.take_needed(&mut self.counts);
1047 let counts = UInt64Array::new(counts.into(), nulls.clone()); let sums = emit_to.take_needed(&mut self.sums);
1050 let sums = PrimitiveArray::<S>::new(sums.into(), nulls) .with_data_type(self.sum_data_type.clone());
1052
1053 Ok(vec![
1054 Arc::new(counts) as ArrayRef,
1055 Arc::new(sums) as ArrayRef,
1056 ])
1057 }
1058
1059 fn merge_batch(
1060 &mut self,
1061 values: &[ArrayRef],
1062 group_indices: &[usize],
1063 total_num_groups: usize,
1064 ) -> Result<()> {
1065 assert_eq!(values.len(), 2, "two arguments to merge_batch");
1066 let partial_counts = values[0].as_primitive::<UInt64Type>();
1068 let partial_sums = values[1].as_primitive::<S>();
1069 self.counts.resize(total_num_groups, 0);
1071 self.null_state.accumulate(
1072 group_indices,
1073 partial_counts,
1074 None,
1075 total_num_groups,
1076 |group_index, partial_count| {
1077 let count = unsafe { self.counts.get_unchecked_mut(group_index) };
1079 *count += partial_count;
1080 },
1081 );
1082
1083 self.sums.resize(total_num_groups, S::default_value());
1085 self.null_state.accumulate(
1086 group_indices,
1087 partial_sums,
1088 None,
1089 total_num_groups,
1090 |group_index, new_value: <S as ArrowPrimitiveType>::Native| {
1091 let sum = unsafe { self.sums.get_unchecked_mut(group_index) };
1093 *sum = sum.add_wrapping(new_value);
1094 },
1095 );
1096
1097 Ok(())
1098 }
1099
1100 fn convert_to_state(
1101 &self,
1102 values: &[ArrayRef],
1103 opt_filter: Option<&BooleanArray>,
1104 ) -> Result<Vec<ArrayRef>> {
1105 let sums = match values[0].as_any().downcast_ref::<PrimitiveArray<S>>() {
1112 Some(sums) => sums.clone().with_data_type(self.sum_data_type.clone()),
1113 None => {
1114 let values = values[0].as_primitive::<I>();
1115 let sums: Vec<S::Native> = values
1118 .values()
1119 .iter()
1120 .map(|value| (*value).into())
1121 .collect();
1122 PrimitiveArray::<S>::new(sums.into(), values.nulls().cloned())
1123 .with_data_type(self.sum_data_type.clone())
1124 }
1125 };
1126 let counts = UInt64Array::from_value(1, sums.len());
1127
1128 let nulls = filtered_null_mask(opt_filter, &sums);
1129
1130 let counts = set_nulls(counts, nulls.clone());
1132 let sums = set_nulls(sums, nulls);
1133
1134 Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)])
1135 }
1136 fn size(&self) -> usize {
1137 self.counts.capacity() * size_of::<u64>()
1139 + self.sums.capacity() * size_of::<S::Native>()
1140 + size_of::<Vec<u64>>()
1142 + size_of::<Vec<S::Native>>()
1143 + self.null_state.size()
1145 }
1146}
1147
1148#[cfg(test)]
1149mod tests {
1150 use super::*;
1151 use arrow::array::{
1152 Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array,
1153 DurationSecondArray, Float64Array,
1154 };
1155 use arrow::datatypes::{Schema, i256};
1156
1157 struct AvgCase {
1158 name: &'static str,
1159 values: ArrayRef,
1160 return_type: DataType,
1161 sum_type: DataType,
1162 expected: ScalarValue,
1163 }
1164
1165 fn with_avg_args<R>(
1166 input_type: &DataType,
1167 return_type: &DataType,
1168 f: impl FnOnce(AccumulatorArgs) -> R,
1169 ) -> R {
1170 let schema = Schema::empty();
1171 let expr_field = Arc::new(Field::new("a", input_type.clone(), true));
1172 let return_field = Arc::new(Field::new("avg", return_type.clone(), true));
1173
1174 f(AccumulatorArgs {
1175 return_field,
1176 schema: &schema,
1177 expr_fields: &[expr_field],
1178 ignore_nulls: false,
1179 order_bys: &[],
1180 is_distinct: false,
1181 name: "avg",
1182 is_reversed: false,
1183 exprs: &[],
1184 })
1185 }
1186
1187 fn avg_groups_accumulator(
1188 input_type: &DataType,
1189 return_type: &DataType,
1190 ) -> Result<Box<dyn GroupsAccumulator>> {
1191 with_avg_args(input_type, return_type, |args| {
1192 Avg::new().create_groups_accumulator(args)
1193 })
1194 }
1195
1196 fn avg_accumulator(
1197 input_type: &DataType,
1198 return_type: &DataType,
1199 ) -> Result<Box<dyn Accumulator>> {
1200 with_avg_args(input_type, return_type, |args| Avg::new().accumulator(args))
1201 }
1202
1203 fn avg_state_fields(
1204 input_type: &DataType,
1205 return_type: &DataType,
1206 ) -> Result<Vec<FieldRef>> {
1207 let input_field = Arc::new(Field::new("a", input_type.clone(), true));
1208 let return_field = Arc::new(Field::new("avg", return_type.clone(), true));
1209
1210 Avg::new().state_fields(StateFieldsArgs {
1211 name: "avg",
1212 input_fields: &[input_field],
1213 return_field,
1214 ordering_fields: &[],
1215 is_distinct: false,
1216 })
1217 }
1218
1219 fn avg_cases() -> Result<Vec<AvgCase>> {
1220 const ROWS: usize = 21_476;
1221 const DECIMAL32_VALUE: i32 = 99_999;
1222 const DECIMAL64_ROWS: usize = 92_235;
1223 const DECIMAL64_VALUE: i64 = 99_999_999_999_999;
1224 const DECIMAL128_ROWS: usize = 21_476;
1225 const DECIMAL128_VALUE: i128 = 9_999_999_999_999_999_999_999_999_999_999_999;
1226
1227 Ok(vec![
1228 AvgCase {
1229 name: "float64",
1230 values: Arc::new(Float64Array::from(vec![10.0, 20.0])),
1231 return_type: DataType::Float64,
1232 sum_type: DataType::Float64,
1233 expected: ScalarValue::Float64(Some(15.0)),
1234 },
1235 AvgCase {
1236 name: "decimal32",
1237 values: Arc::new(
1238 Decimal32Array::from(vec![Some(DECIMAL32_VALUE); ROWS])
1239 .with_precision_and_scale(5, 0)?,
1240 ),
1241 return_type: DataType::Decimal32(9, 4),
1242 sum_type: DataType::Decimal64(18, 0),
1243 expected: ScalarValue::Decimal32(Some(DECIMAL32_VALUE * 10_000), 9, 4),
1244 },
1245 AvgCase {
1246 name: "decimal64",
1247 values: Arc::new(
1248 Decimal64Array::from(vec![Some(DECIMAL64_VALUE); DECIMAL64_ROWS])
1249 .with_precision_and_scale(14, 0)?,
1250 ),
1251 return_type: DataType::Decimal64(18, 4),
1252 sum_type: DataType::Decimal128(38, 0),
1253 expected: ScalarValue::Decimal64(Some(DECIMAL64_VALUE * 10_000), 18, 4),
1254 },
1255 AvgCase {
1256 name: "decimal128",
1257 values: Arc::new(
1258 Decimal128Array::from(vec![Some(DECIMAL128_VALUE); DECIMAL128_ROWS])
1259 .with_precision_and_scale(34, 0)?,
1260 ),
1261 return_type: DataType::Decimal128(38, 4),
1262 sum_type: DataType::Decimal256(76, 0),
1263 expected: ScalarValue::Decimal128(Some(DECIMAL128_VALUE * 10_000), 38, 4),
1264 },
1265 AvgCase {
1266 name: "decimal256",
1267 values: Arc::new(
1268 Decimal256Array::from(vec![i256::from_i128(10), i256::from_i128(20)])
1269 .with_precision_and_scale(50, 0)?,
1270 ),
1271 return_type: DataType::Decimal256(54, 4),
1272 sum_type: DataType::Decimal256(76, 0),
1273 expected: ScalarValue::Decimal256(Some(i256::from_i128(150_000)), 54, 4),
1274 },
1275 AvgCase {
1278 name: "decimal128_with_headroom",
1279 values: Arc::new(
1280 Decimal128Array::from(vec![100_000, 200_000])
1281 .with_precision_and_scale(20, 4)?,
1282 ),
1283 return_type: DataType::Decimal128(24, 8),
1284 sum_type: DataType::Decimal128(38, 4),
1285 expected: ScalarValue::Decimal128(Some(1_500_000_000), 24, 8),
1286 },
1287 AvgCase {
1290 name: "decimal32_max_precision",
1291 values: Arc::new(
1292 Decimal32Array::from(vec![10, 20]).with_precision_and_scale(9, 0)?,
1293 ),
1294 return_type: DataType::Decimal32(9, 4),
1295 sum_type: DataType::Decimal128(38, 0),
1296 expected: ScalarValue::Decimal32(Some(150_000), 9, 4),
1297 },
1298 AvgCase {
1301 name: "duration_second",
1302 values: Arc::new(DurationSecondArray::from(vec![10, 20])),
1303 return_type: DataType::Duration(TimeUnit::Second),
1304 sum_type: DataType::Duration(TimeUnit::Second),
1305 expected: ScalarValue::DurationSecond(Some(15)),
1306 },
1307 ])
1308 }
1309
1310 #[test]
1311 fn avg_accumulator_evaluate_and_state_types() -> Result<()> {
1312 for case in avg_cases()? {
1313 let input_type = case.values.data_type();
1314 let state_fields = avg_state_fields(input_type, &case.return_type)?;
1315 let mut acc = avg_accumulator(input_type, &case.return_type)?;
1316 acc.update_batch(std::slice::from_ref(&case.values))?;
1317
1318 let state = acc.state()?;
1319 assert_eq!(
1320 &state[0].data_type(),
1321 state_fields[0].data_type(),
1322 "{}",
1323 case.name
1324 );
1325 assert_eq!(
1326 &state[1].data_type(),
1327 state_fields[1].data_type(),
1328 "{}",
1329 case.name
1330 );
1331 assert_eq!(acc.evaluate()?, case.expected, "{}", case.name);
1332 }
1333
1334 Ok(())
1335 }
1336
1337 #[test]
1338 fn avg_groups_state_types_match_state_fields() -> Result<()> {
1339 for case in avg_cases()? {
1340 let input_type = case.values.data_type();
1341 let state_fields = avg_state_fields(input_type, &case.return_type)?;
1342 let acc = avg_groups_accumulator(input_type, &case.return_type)?;
1343 let state = acc.convert_to_state(std::slice::from_ref(&case.values), None)?;
1344
1345 assert_eq!(
1346 state_fields[0].data_type(),
1347 &DataType::UInt64,
1348 "{}",
1349 case.name
1350 );
1351 assert_eq!(state_fields[1].data_type(), &case.sum_type, "{}", case.name);
1352 assert_eq!(state[0].data_type(), &DataType::UInt64, "{}", case.name);
1353 assert_eq!(state[1].data_type(), &case.sum_type, "{}", case.name);
1354 }
1355
1356 Ok(())
1357 }
1358
1359 #[test]
1360 fn avg_groups_convert_to_state_roundtrip() -> Result<()> {
1361 for case in avg_cases()? {
1362 let input_type = case.values.data_type();
1363 let partial = avg_groups_accumulator(input_type, &case.return_type)?;
1364 let mut final_acc = avg_groups_accumulator(input_type, &case.return_type)?;
1365 let state =
1366 partial.convert_to_state(std::slice::from_ref(&case.values), None)?;
1367 final_acc.merge_batch(&state, &vec![0; case.values.len()], 1)?;
1368
1369 let result = final_acc.evaluate(EmitTo::All)?;
1370 assert_eq!(result.data_type(), &case.return_type, "{}", case.name);
1371 assert_eq!(
1372 ScalarValue::try_from_array(result.as_ref(), 0)?,
1373 case.expected,
1374 "{}",
1375 case.name
1376 );
1377 }
1378
1379 Ok(())
1380 }
1381
1382 #[test]
1385 fn avg_errors_when_average_exceeds_output_precision() -> Result<()> {
1386 let values: ArrayRef = Arc::new(
1387 Decimal32Array::from(vec![999_999_999]).with_precision_and_scale(9, 0)?,
1388 );
1389 let return_type = DataType::Decimal32(9, 4);
1390 let mut acc = avg_accumulator(values.data_type(), &return_type)?;
1391
1392 acc.update_batch(&[values])?;
1393 assert!(acc.evaluate().is_err());
1394
1395 Ok(())
1396 }
1397}