Skip to content

Commit 31335b4

Browse files
committed
more plumbing
1 parent 363c3ee commit 31335b4

File tree

1 file changed

+57
-1
lines changed

1 file changed

+57
-1
lines changed

datafusion/physical-expr/src/aggregate/average.rs

+57-1
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ use crate::aggregate::sum::sum_batch;
2929
use crate::aggregate::utils::calculate_result_decimal_for_avg;
3030
use crate::aggregate::utils::down_cast_any_ref;
3131
use crate::expressions::format_state_name;
32-
use crate::{AggregateExpr, PhysicalExpr};
32+
use crate::{AggregateExpr, GroupsAccumulator, PhysicalExpr};
3333
use arrow::compute;
3434
use arrow::datatypes::DataType;
3535
use arrow::{
@@ -155,6 +155,13 @@ impl AggregateExpr for Avg {
155155
&self.rt_data_type,
156156
)?))
157157
}
158+
159+
fn create_groups_accumulator(&self) -> Result<Box<dyn GroupsAccumulator>> {
160+
Ok(Box::new(AvgGroupsAccumulator::new(
161+
&self.sum_data_type,
162+
&self.rt_data_type,
163+
)))
164+
}
158165
}
159166

160167
impl PartialEq<dyn Any> for Avg {
@@ -383,6 +390,55 @@ impl RowAccumulator for AvgRowAccumulator {
383390
}
384391
}
385392

393+
/// An accumulator to compute the average
394+
#[derive(Debug)]
395+
pub struct AvgGroupsAccumulator {
396+
sum_data_type: DataType,
397+
return_data_type: DataType,
398+
//count: u64,
399+
}
400+
401+
impl AvgGroupsAccumulator {
402+
pub fn new(sum_datatype: &DataType, return_data_type: &DataType) -> Self {
403+
Self {
404+
sum_data_type: sum_datatype.clone(),
405+
return_data_type: return_data_type.clone(),
406+
}
407+
}
408+
}
409+
410+
impl GroupsAccumulator for AvgGroupsAccumulator {
411+
fn update_batch(
412+
&mut self,
413+
values: &[ArrayRef],
414+
group_indicies: &[usize],
415+
opt_filter: Option<&arrow_array::BooleanArray>,
416+
) -> Result<usize> {
417+
todo!()
418+
}
419+
420+
fn evaluate(&mut self) -> Result<ArrayRef> {
421+
todo!()
422+
}
423+
424+
fn state(&mut self) -> Result<Vec<ArrayRef>> {
425+
todo!()
426+
}
427+
428+
fn merge_batch(
429+
&mut self,
430+
values: &[ArrayRef],
431+
group_indicies: &[usize],
432+
opt_filter: Option<&arrow_array::BooleanArray>,
433+
) -> Result<()> {
434+
todo!()
435+
}
436+
437+
fn size(&self) -> usize {
438+
todo!()
439+
}
440+
}
441+
386442
#[cfg(test)]
387443
mod tests {
388444
use super::*;

0 commit comments

Comments
 (0)