first iteration on intern()

Author: Rich-T-kid

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

@@ -41,10 +41,13 @@ pub(crate) use single_group_by::primitive::HashValue;
 use crate::aggregates::{
     group_values::single_group_by::{
         boolean::GroupValuesBoolean, bytes::GroupValuesBytes,
-        bytes_view::GroupValuesBytesView, primitive::GroupValuesPrimitive, dictionary::GroupValuesDictionary,
+        bytes_view::GroupValuesBytesView, dictionary::GroupValuesDictionary,
+        primitive::GroupValuesPrimitive,
     },
     order::GroupOrdering,
 };
+use arrow::array::*;
+use std::sync::Arc;
 
 mod metrics;
 mod null_builder;
@@ -199,11 +202,24 @@ pub fn new_group_values(
             DataType::Boolean => {
                 return Ok(Box::new(GroupValuesBoolean::new()));
             }
-            /*DataType::Dictionary(_, _) => {
-                println!("dictionary type detected, using SingleDictionaryGroupValues");
-                return Ok(Box::new(SingleDictionaryGroupValues::new()));
-
-            }*/
+            DataType::Dictionary(key_type, value_type) => {
+                if supported_single_dictionary_value(value_type) {
+                    println!("dictionary type detected, using GroupValuesDictionary");
+                    return match key_type.as_ref() { // TODO: turn this into a macro
+                        DataType::Int8 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::Int8Type>::new(value_type))),
+                        DataType::Int16 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::Int16Type>::new(value_type))),
+                        DataType::Int32 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::Int32Type>::new(value_type))),
+                        DataType::Int64 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::Int64Type>::new(value_type))),
+                        DataType::UInt8 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::UInt8Type>::new(value_type))),
+                        DataType::UInt16 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::UInt16Type>::new(value_type))),
+                        DataType::UInt32 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::UInt32Type>::new(value_type))),
+                        DataType::UInt64 => Ok(Box::new(GroupValuesDictionary::<arrow::datatypes::UInt64Type>::new(value_type))),
+                        _ => Err(datafusion_common::DataFusionError::NotImplemented(
+                            format!("Unsupported dictionary key type: {:?}", key_type)
+                        )),
+                    };
+                }
+            }
             _ => {}
         }
     }
@@ -219,3 +235,24 @@ pub fn new_group_values(
         Ok(Box::new(GroupValuesRows::try_new(schema)?))
     }
 }
+
+fn supported_single_dictionary_value(t: &DataType) -> bool {
+    matches!(
+        t,
+        DataType::Utf8
+            | DataType::LargeUtf8
+            | DataType::Binary
+            | DataType::LargeBinary
+            | DataType::Utf8View
+            | DataType::BinaryView
+            | DataType::Int8
+            | DataType::Int16
+            | DataType::Int32
+            | DataType::Int64
+            | DataType::UInt8
+            | DataType::UInt16
+            | DataType::UInt32
+            | DataType::UInt64
+    )
+}
+

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

@@ -473,24 +473,21 @@ mod playground {
 
     #[tokio::test]
     async fn test_trivial_group_by_dictionary() -> Result<()> {
+        use crate::aggregates::RecordBatch;
+        use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+        use crate::common::collect;
+        use crate::test::TestMemoryExec;
         use arrow::array::DictionaryArray;
         use arrow::datatypes::{DataType, Field, Schema};
         use datafusion_functions_aggregate::count::count_udaf;
         use datafusion_physical_expr::aggregate::AggregateExprBuilder;
         use datafusion_physical_expr::expressions::col;
-        use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
-        use crate::test::TestMemoryExec;
-        use crate::aggregates::RecordBatch;
-        use crate::common::collect;
 
         // Create schema with dictionary column and value column
         let schema = Arc::new(Schema::new(vec![
             Field::new(
                 "color",
-                DataType::Dictionary(
-                    Box::new(DataType::UInt8),
-                    Box::new(DataType::Utf8),
-                ),
+                DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
                 false,
             ),
             Field::new("amount", DataType::UInt32, false),
@@ -499,11 +496,9 @@ mod playground {
         // Create dictionary array
         let values = StringArray::from(vec!["red", "blue", "green"]);
         let keys = arrow::array::UInt8Array::from(vec![0, 1, 0, 2, 1]);
-        let dict_array: ArrayRef =
-            Arc::new(DictionaryArray::<arrow::datatypes::UInt8Type>::try_new(
-                keys,
-                Arc::new(values),
-            )?);
+        let dict_array: ArrayRef = Arc::new(DictionaryArray::<
+            arrow::datatypes::UInt8Type,
+        >::try_new(keys, Arc::new(values))?);
 
         // Create value column
         let amount_array: ArrayRef =
@@ -514,7 +509,8 @@ mod playground {
             RecordBatch::try_new(Arc::clone(&schema), vec![dict_array, amount_array])?;
 
         // Create in-memory source with the batch
-        let source = TestMemoryExec::try_new(&vec![vec![batch]], Arc::clone(&schema), None)?;
+        let source =
+            TestMemoryExec::try_new(&vec![vec![batch]], Arc::clone(&schema), None)?;
 
         // Create GROUP BY expression
         let group_expr = vec![(col("color", &schema)?, "color".to_string())];
@@ -537,11 +533,9 @@ mod playground {
             Arc::clone(&schema),
         )?;
 
-
         let output =
             collect(aggregate_exec.execute(0, Arc::new(TaskContext::default()))?).await?;
         println!("Output batch: {:#?}", output);
         Ok(())
     }
-
 }