Support for Avg duration

Author: shruti2522

datafusion/functions-aggregate/src/average.rs

@@ -171,7 +171,7 @@ impl AggregateUDFImpl for Avg {
     fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
         matches!(
             args.return_type,
-            DataType::Float64 | DataType::Decimal128(_, _)
+            DataType::Float64 | DataType::Decimal128(_, _) | DataType::Duration(_)
         )
     }
 
@@ -232,6 +232,49 @@ impl AggregateUDFImpl for Avg {
                 )))
             }
 
+            (Duration(time_unit), Duration(_)) => match time_unit {
+                TimeUnit::Second => {
+                    let avg_fn = move |sum: i64, count: u64| Ok(sum / count as i64);
+                    Ok(Box::new(
+                        AvgGroupsAccumulator::<DurationSecondType, _>::new(
+                            &data_type,
+                            args.return_type,
+                            avg_fn,
+                        ),
+                    ))
+                }
+                TimeUnit::Millisecond => {
+                    let avg_fn = move |sum: i64, count: u64| Ok(sum / count as i64);
+                    Ok(Box::new(
+                        AvgGroupsAccumulator::<DurationMillisecondType, _>::new(
+                            &data_type,
+                            args.return_type,
+                            avg_fn,
+                        ),
+                    ))
+                }
+                TimeUnit::Microsecond => {
+                    let avg_fn = move |sum: i64, count: u64| Ok(sum / count as i64);
+                    Ok(Box::new(
+                        AvgGroupsAccumulator::<DurationMicrosecondType, _>::new(
+                            &data_type,
+                            args.return_type,
+                            avg_fn,
+                        ),
+                    ))
+                }
+                TimeUnit::Nanosecond => {
+                    let avg_fn = move |sum: i64, count: u64| Ok(sum / count as i64);
+                    Ok(Box::new(
+                        AvgGroupsAccumulator::<DurationNanosecondType, _>::new(
+                            &data_type,
+                            args.return_type,
+                            avg_fn,
+                        ),
+                    ))
+                }
+            },
+
             _ => not_impl_err!(
                 "AvgGroupsAccumulator for ({} --> {})",
                 &data_type,

datafusion/sqllogictest/test_files/aggregate.slt

@@ -4969,6 +4969,34 @@ select count(distinct column1), count(distinct column2) from dict_test group by
 statement ok
 drop table dict_test;
 
+# avg_duartion
+
+statement ok
+create table d as values 
+  (arrow_cast(1, 'Duration(Second)'), arrow_cast(2, 'Duration(Millisecond)'), arrow_cast(3, 'Duration(Microsecond)'), arrow_cast(4, 'Duration(Nanosecond)'), 1),
+  (arrow_cast(11, 'Duration(Second)'), arrow_cast(22, 'Duration(Millisecond)'), arrow_cast(33, 'Duration(Microsecond)'), arrow_cast(44, 'Duration(Nanosecond)'), 1);
+
+query ????
+SELECT avg(column1), avg(column2), avg(column3), avg(column4) FROM d;
+----
+0 days 0 hours 0 mins 6 secs 0 days 0 hours 0 mins 0.012 secs 0 days 0 hours 0 mins 0.000018 secs 0 days 0 hours 0 mins 0.000000024 secs
+
+query ????I
+SELECT avg(column1), avg(column2), avg(column3), avg(column4), column5 FROM d GROUP BY column5;
+----
+0 days 0 hours 0 mins 6 secs 0 days 0 hours 0 mins 0.012 secs 0 days 0 hours 0 mins 0.000018 secs 0 days 0 hours 0 mins 0.000000024 secs 1
+
+INSERT INTO d VALUES
+  (arrow_cast(3, 'Duration(Second)'), arrow_cast(1, 'Duration(Millisecond)'), arrow_cast(7, 'Duration(Microsecond)'), arrow_cast(2, 'Duration(Nanosecond)'), 1),
+  (arrow_cast(0, 'Duration(Second)'), arrow_cast(9, 'Duration(Millisecond)'), arrow_cast(5, 'Duration(Microsecond)'), arrow_cast(8, 'Duration(Nanosecond)'), 1);
+
+query ????I
+SELECT avg(column1), avg(column2), avg(column3), avg(column4), column5 FROM d GROUP BY column5 ORDER BY column5;
+----
+0 days 0 hours 0 mins 11 secs 0 days 0 hours 0 mins 0.022 secs 0 days 0 hours 0 mins 0.000033 secs 0 days 0 hours 0 mins 0.000000044 secs 1
+
+statement ok
+drop table d;
 
 # Prepare the table with dictionary values for testing
 statement ok