Researching code flow

Author: Rich-T-kid

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

@@ -41,7 +41,7 @@ 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,
+        bytes_view::GroupValuesBytesView, primitive::GroupValuesPrimitive, dictionary::GroupValuesDictionary,
     },
     order::GroupOrdering,
 };
@@ -137,6 +137,9 @@ pub fn new_group_values(
 ) -> Result<Box<dyn GroupValues>> {
     if schema.fields.len() == 1 {
         let d = schema.fields[0].data_type();
+        println!(
+            "[should be dictionary encoded] single column group by with data type: {d:#?}"
+        );
 
         macro_rules! downcast_helper {
             ($t:ty, $d:ident) => {
@@ -196,6 +199,11 @@ 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()));
+
+            }*/
             _ => {}
         }
     }
@@ -207,6 +215,7 @@ pub fn new_group_values(
             Ok(Box::new(GroupValuesColumn::<true>::try_new(schema)?))
         }
     } else {
+        // TODO:  add specialized implementation for dictionary encoding columns for 2+ group by columns case
         Ok(Box::new(GroupValuesRows::try_new(schema)?))
     }
 }

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

@@ -23,9 +23,9 @@ use arrow::array::{
 use arrow::compute::cast;
 use arrow::datatypes::{DataType, SchemaRef};
 use arrow::row::{RowConverter, Rows, SortField};
-use datafusion_common::Result;
 use datafusion_common::hash_utils::RandomState;
 use datafusion_common::hash_utils::create_hashes;
+use datafusion_common::{Result, internal_err};
 use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt};
 use datafusion_expr::EmitTo;
 use hashbrown::hash_table::HashTable;
@@ -324,3 +324,224 @@ fn dictionary_encode_if_necessary(
         (_, _) => Ok(Arc::<dyn Array>::clone(array)),
     }
 }
+
+mod playground {
+    use std::ops::Index;
+
+    use arrow::array::{AsArray, Datum};
+    use datafusion_execution::TaskContext;
+
+    use crate::{ExecutionPlan, test::TestMemoryExec};
+
+    use super::*;
+    use arrow::array::{Array, ArrayRef, StringArray};
+    use std::string::String;
+    struct TrivialGroupBy {
+        seen_strings: Vec<String>,
+        cur_size: usize,
+    }
+    impl GroupValues for TrivialGroupBy {
+        // trivial apprach assume theres only one columns
+        fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+            let n_rows = cols[0].len();
+            let column_count = cols.len();
+            // iterate all the rows, for each row, concate each value into a final string
+            for row_idx in 0..n_rows {
+                // grab the underlying value; assume its a string
+                let mut cur_str = String::from("");
+                for col_idx in 0..column_count {
+                    let array = cols.get(col_idx).unwrap();
+                    let string_array =
+                        array.as_any().downcast_ref::<StringArray>().unwrap();
+                    if string_array.is_valid(row_idx) {
+                        let value = string_array.value(row_idx);
+                        cur_str.push_str(value);
+                    }
+                }
+                let idx = if let Some(i) =
+                    self.seen_strings.iter().position(|x| x == &cur_str)
+                {
+                    i
+                } else {
+                    self.seen_strings.push(cur_str);
+                    self.seen_strings.len() - 1
+                };
+                groups.push(idx);
+            }
+            println!("{:?}", self.seen_strings);
+            Ok(())
+        }
+        fn len(&self) -> usize {
+            self.seen_strings.len()
+        }
+        fn is_empty(&self) -> bool {
+            self.seen_strings.is_empty()
+        }
+        fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+            let strings_to_emit = match emit_to {
+                EmitTo::All => {
+                    // take all groups, clear internal state
+                    std::mem::take(&mut self.seen_strings)
+                }
+                EmitTo::First(n) => {
+                    // take only the first n groups
+                    // drain removes them from seen_strings
+                    self.seen_strings.drain(..n).collect()
+                }
+            };
+
+            // convert our Vec<String> back into an Arrow StringArray
+            let array: ArrayRef = Arc::new(StringArray::from(
+                strings_to_emit
+                    .iter()
+                    .map(|s| s.as_str())
+                    .collect::<Vec<_>>(),
+            ));
+
+            // return one array per GROUP BY column
+            // since we're trivial and only support one column, just return one
+            Ok(vec![array])
+        }
+
+        fn size(&self) -> usize {
+            self.cur_size
+        }
+        fn clear_shrink(&mut self, num_rows: usize) {
+            let _ = num_rows;
+        }
+    }
+
+    #[test]
+    fn test_trivial_group_by_single_column() {
+        // Test grouping on a single string column
+        let strings = vec!["apple", "banana", "apple", "cherry", "banana"];
+        let array: ArrayRef = Arc::new(StringArray::from(strings));
+
+        let mut group_by = TrivialGroupBy {
+            seen_strings: Vec::new(),
+            cur_size: 0,
+        };
+
+        // Intern the group keys
+        let mut groups = Vec::new();
+        group_by.intern(&[array], &mut groups).unwrap();
+
+        // Should have assigned group ids: 0, 1, 0, 2, 1
+        assert_eq!(groups, vec![0, 1, 0, 2, 1]);
+        assert_eq!(group_by.len(), 3); // apple, banana, cherry
+
+        // Emit all groups
+        let emitted = group_by.emit(EmitTo::All).unwrap();
+        assert_eq!(emitted.len(), 1); // One column
+        let emitted_array = emitted[0].as_any().downcast_ref::<StringArray>().unwrap();
+        assert_eq!(emitted_array.value(0), "apple");
+        assert_eq!(emitted_array.value(1), "banana");
+        assert_eq!(emitted_array.value(2), "cherry");
+    }
+
+    #[test]
+    fn test_trivial_group_by_two_columns() {
+        // Test grouping on two string columns
+        let col1 = vec!["a", "a", "b", "a", "b"];
+        let col2 = vec!["x", "y", "x", "x", "y"];
+
+        let array1: ArrayRef = Arc::new(StringArray::from(col1));
+        let array2: ArrayRef = Arc::new(StringArray::from(col2));
+
+        let mut group_by = TrivialGroupBy {
+            seen_strings: Vec::new(),
+            cur_size: 0,
+        };
+
+        // Intern: concatenates ("a" + "x"), ("a" + "y"), ("b" + "x"), etc.
+        let mut groups = Vec::new();
+        group_by.intern(&[array1, array2], &mut groups).unwrap();
+
+        // Should have 4 distinct groups: "ax", "ay", "bx", "by"
+        assert_eq!(group_by.len(), 4);
+        assert_eq!(groups, vec![0, 1, 2, 0, 3]); // group ids assigned
+
+        // Emit all groups
+        let emitted = group_by.emit(EmitTo::All).unwrap();
+        assert_eq!(emitted.len(), 1); // One output column (concatenated strings)
+        let emitted_array = emitted[0].as_any().downcast_ref::<StringArray>().unwrap();
+        assert_eq!(emitted_array.value(0), "ax");
+        assert_eq!(emitted_array.value(1), "ay");
+        assert_eq!(emitted_array.value(2), "bx");
+        assert_eq!(emitted_array.value(3), "by");
+    }
+
+    #[tokio::test]
+    async fn test_trivial_group_by_dictionary() -> Result<()> {
+        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),
+                ),
+                false,
+            ),
+            Field::new("amount", DataType::UInt32, false),
+        ]));
+
+        // 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),
+            )?);
+
+        // Create value column
+        let amount_array: ArrayRef =
+            Arc::new(arrow::array::UInt32Array::from(vec![1, 2, 3, 4, 5]));
+
+        // Create batch
+        let batch =
+            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)?;
+
+        // Create GROUP BY expression
+        let group_expr = vec![(col("color", &schema)?, "color".to_string())];
+
+        // Create COUNT(amount) aggregate expression
+        let aggr_expr = vec![Arc::new(
+            AggregateExprBuilder::new(count_udaf(), vec![col("amount", &schema)?])
+                .schema(Arc::clone(&schema))
+                .alias("count_amount")
+                .build()?,
+        )];
+
+        // Create AggregateExec
+        let aggregate_exec = AggregateExec::try_new(
+            AggregateMode::SinglePartitioned,
+            PhysicalGroupBy::new_single(group_expr),
+            aggr_expr,
+            vec![None],
+            Arc::new(source),
+            Arc::clone(&schema),
+        )?;
+
+
+        let output =
+            collect(aggregate_exec.execute(0, Arc::new(TaskContext::default()))?).await?;
+        println!("Output batch: {:#?}", output);
+        Ok(())
+    }
+
+}