support_ansi_avg_wip

Author: coderfender

native/core/src/execution/planner.rs

@@ -113,10 +113,10 @@ use datafusion_comet_proto::{
 };
 use datafusion_comet_spark_expr::monotonically_increasing_id::MonotonicallyIncreasingId;
 use datafusion_comet_spark_expr::{
-    ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct,
-    GetArrayStructFields, GetStructField, IfExpr, ListExtract, NormalizeNaNAndZero, RLike,
-    RandExpr, RandnExpr, SparkCastOptions, Stddev, SubstringExpr, SumDecimal, TimestampTruncExpr,
-    ToJson, UnboundColumn, Variance, AvgInt
+    ArrayInsert, Avg, AvgDecimal, AvgInt, Cast, CheckOverflow, Correlation, Covariance,
+    CreateNamedStruct, GetArrayStructFields, GetStructField, IfExpr, ListExtract,
+    NormalizeNaNAndZero, RLike, RandExpr, RandnExpr, SparkCastOptions, Stddev, SubstringExpr,
+    SumDecimal, TimestampTruncExpr, ToJson, UnboundColumn, Variance,
 };
 use itertools::Itertools;
 use jni::objects::GlobalRef;
@@ -1854,8 +1854,11 @@ impl PhysicalPlanner {
                 let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap());
                 let input_datatype = to_arrow_datatype(expr.sum_datatype.as_ref().unwrap());
                 let builder = match datatype {
-                    
-                    DataType::Int8 | DataType::UInt8 | DataType::Int16 | DataType::UInt16 | DataType::Int32 => {
+                    DataType::Int8
+                    | DataType::UInt8
+                    | DataType::Int16
+                    | DataType::UInt16
+                    | DataType::Int32 => {
                         let func =
                             AggregateUDF::new_from_impl(AvgInt::new(datatype, input_datatype));
                         AggregateExprBuilder::new(Arc::new(func), vec![child])

native/spark-expr/src/agg_funcs/avg_int.rs

@@ -15,15 +15,17 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::any::Any;
+use crate::{AvgDecimal, EvalMode};
 use arrow::array::{ArrayRef, BooleanArray};
 use arrow::datatypes::{DataType, FieldRef};
 use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
-use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature};
 use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion::logical_expr::type_coercion::aggregates::avg_return_type;
 use datafusion::logical_expr::Volatility::Immutable;
-use crate::{AvgDecimal, EvalMode};
+use datafusion::logical_expr::{
+    Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature,
+};
+use std::any::Any;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct AvgInt {
@@ -34,13 +36,13 @@ pub struct AvgInt {
 impl AvgInt {
     pub fn try_new(data_type: DataType, eval_mode: EvalMode) -> DFResult<Self> {
         match data_type {
-            DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
-                Ok(Self {
+            DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => Ok(Self {
                 signature: Signature::user_defined(Immutable),
-                eval_mode
-                })
-            },
-            _ => {Err(DataFusionError::Internal("inalid data type for AvgInt".to_string()))}
+                eval_mode,
+            }),
+            _ => Err(DataFusionError::Internal(
+                "inalid data type for AvgInt".to_string(),
+            )),
         }
     }
 }
@@ -58,7 +60,7 @@ impl AggregateUDFImpl for AvgInt {
         ReversedUDAF::Identical
     }
 
-    fn signature(&self) -> &Signature  {
+    fn signature(&self) -> &Signature {
         &self.signature
     }
 
@@ -70,7 +72,10 @@ impl AggregateUDFImpl for AvgInt {
         true
     }
 
-    fn accumulator(&self, acc_args: AccumulatorArgs) -> datafusion::common::Result<Box<dyn Accumulator>> {
+    fn accumulator(
+        &self,
+        acc_args: AccumulatorArgs,
+    ) -> datafusion::common::Result<Box<dyn Accumulator>> {
         todo!()
     }
 
@@ -82,7 +87,10 @@ impl AggregateUDFImpl for AvgInt {
         false
     }
 
-    fn create_groups_accumulator(&self, _args: AccumulatorArgs) -> datafusion::common::Result<Box<dyn GroupsAccumulator>> {
+    fn create_groups_accumulator(
+        &self,
+        _args: AccumulatorArgs,
+    ) -> datafusion::common::Result<Box<dyn GroupsAccumulator>> {
         Ok(Box::new(AvgIntGroupsAccumulator::new(self.eval_mode)))
     }
 
@@ -91,37 +99,36 @@ impl AggregateUDFImpl for AvgInt {
     }
 }
 
-struct AvgIntegerAccumulator{
+struct AvgIntegerAccumulator {
     sum: Option<i64>,
     count: u64,
     eval_mode: EvalMode,
 }
 
 impl AvgIntegerAccumulator {
     fn new(eval_mode: EvalMode) -> Self {
-        Self{
-            sum : Some(0),
+        Self {
+            sum: Some(0),
             count: 0,
-            eval_mode
+            eval_mode,
         }
     }
 }
 
-impl Accumulator for AvgIntegerAccumulator {
-    
-}
+impl Accumulator for AvgIntegerAccumulator {}
 
-struct AvgIntGroupsAccumulator {
-
-}
-
-impl AvgIntGroupsAccumulator {
-
-}
+struct AvgIntGroupsAccumulator {}
 
+impl AvgIntGroupsAccumulator {}
 
 impl GroupsAccumulator for AvgIntGroupsAccumulator {
-    fn update_batch(&mut self, values: &[ArrayRef], group_indices: &[usize], opt_filter: Option<&BooleanArray>, total_num_groups: usize) -> datafusion::common::Result<()> {
+    fn update_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> datafusion::common::Result<()> {
         todo!()
     }
 
@@ -133,11 +140,17 @@ impl GroupsAccumulator for AvgIntGroupsAccumulator {
         todo!()
     }
 
-    fn merge_batch(&mut self, values: &[ArrayRef], group_indices: &[usize], opt_filter: Option<&BooleanArray>, total_num_groups: usize) -> datafusion::common::Result<()> {
+    fn merge_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> datafusion::common::Result<()> {
         todo!()
     }
 
     fn size(&self) -> usize {
         todo!()
     }
-}
+}

native/spark-expr/src/agg_funcs/mod.rs

@@ -17,18 +17,18 @@
 
 mod avg;
 mod avg_decimal;
+mod avg_int;
 mod correlation;
 mod covariance;
 mod stddev;
 mod sum_decimal;
 mod variance;
-mod avg_int;
 
 pub use avg::Avg;
 pub use avg_decimal::AvgDecimal;
+pub use avg_int::AvgInt;
 pub use correlation::Correlation;
 pub use covariance::Covariance;
 pub use stddev::Stddev;
 pub use sum_decimal::SumDecimal;
 pub use variance::Variance;
-pub use avg_int::AvgInt;