revised-version-1

Author: Rich-T-kid

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -23,7 +23,11 @@ use arrow::array::types::{
     TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
 };
 use arrow::array::{ArrayRef, downcast_primitive};
-use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
+
+use arrow::datatypes::{
+    DataType, Int8Type, Int16Type, Int32Type, Int64Type, SchemaRef, TimeUnit, UInt8Type,
+    UInt16Type, UInt32Type, UInt64Type,
+};
 use datafusion_common::Result;
 
 use datafusion_expr::EmitTo;
@@ -36,7 +40,9 @@ use datafusion_physical_expr::binary_map::OutputType;
 use multi_group_by::GroupValuesColumn;
 use row::GroupValuesRows;
 
-pub(crate) use single_group_by::primitive::HashValue;
+pub(crate) use single_group_by::{
+    dictionary::GroupValuesDictionary, primitive::HashValue,
+};
 
 use crate::aggregates::{
     group_values::single_group_by::{
@@ -196,6 +202,45 @@ pub fn new_group_values(
             DataType::Boolean => {
                 return Ok(Box::new(GroupValuesBoolean::new()));
             }
+            DataType::Dictionary(key_type, value_type) => {
+                if supported_single_dictionary_value(value_type) {
+                    return match key_type.as_ref() {
+                        DataType::Int8 => Ok(Box::new(
+                            GroupValuesDictionary::<Int8Type>::new(value_type),
+                        )),
+                        DataType::Int16 => Ok(Box::new(
+                            GroupValuesDictionary::<Int16Type>::new(value_type),
+                        )),
+                        DataType::Int32 => Ok(Box::new(
+                            GroupValuesDictionary::<Int32Type>::new(value_type),
+                        )),
+                        DataType::Int64 => Ok(Box::new(
+                            GroupValuesDictionary::<Int64Type>::new(value_type),
+                        )),
+                        DataType::UInt8 => Ok(Box::new(
+                            GroupValuesDictionary::<UInt8Type>::new(value_type),
+                        )),
+                        DataType::UInt16 => {
+                            Ok(Box::new(GroupValuesDictionary::<UInt16Type>::new(
+                                value_type,
+                            )))
+                        }
+                        DataType::UInt32 => {
+                            Ok(Box::new(GroupValuesDictionary::<UInt32Type>::new(
+                                value_type,
+                            )))
+                        }
+                        DataType::UInt64 => {
+                            Ok(Box::new(GroupValuesDictionary::<UInt64Type>::new(
+                                value_type,
+                            )))
+                        }
+                        _ => Err(datafusion_common::DataFusionError::NotImplemented(
+                            format!("Unsupported dictionary key type: {key_type:?}"),
+                        )),
+                    };
+                }
+            }
             _ => {}
         }
     }
@@ -210,3 +255,18 @@ pub fn new_group_values(
         Ok(Box::new(GroupValuesRows::try_new(schema)?))
     }
 }
+
+fn supported_single_dictionary_value(t: &DataType) -> bool {
+    match t {
+        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => true,
+        DataType::List(field)
+            if matches!(
+                field.data_type(),
+                DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8
+            ) =>
+        {
+            true
+        }
+        _ => false,
+    }
+}