all test pass, first iteration done| need to run benchmarks

Author: Rich-T-kid

datafusion/physical-plan/src/aggregates/group_values/single_group_by/dictionary.rs

@@ -17,8 +17,9 @@
 
 use crate::aggregates::group_values::GroupValues;
 use arrow::array::{
-    Array, ArrayBuilder, ArrayRef, DictionaryArray, Int8Builder, Int16Builder, Int32Builder, Int64Builder, LargeStringBuilder, StringArray, StringBuilder, StringViewBuilder, UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder
+    Array, ArrayBuilder, ArrayRef, DictionaryArray, Int8Builder, Int16Builder, Int32Builder, Int64Builder, LargeStringBuilder, Scalar, StringArray, StringBuilder, StringViewBuilder, UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder
 };
+use std::mem;
 use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType};
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::EmitTo;
@@ -76,8 +77,10 @@ impl<K:ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
 impl<K:ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K>  {
     // not really sure how to return the size of strings and binary values so this is a best effort approach
     fn size(&self) -> usize {
-        0
-        }
+        let arr_size = element_size(&self.value_dt) * self.unique_dict_value_mapping.len();
+        let dict_size = self.unique_dict_value_mapping.len() * size_of::<(ScalarValue, usize)>() + 100 /* rough estimate for hashmap overhead */; // rough estimate for hashmap overhead
+        arr_size + dict_size
+    }
     fn len(&self) -> usize {
         self.unique_dict_value_mapping.len()
     }
@@ -92,7 +95,6 @@ impl<K:ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K>
         }
         let array = cols[0].clone();
         groups.clear(); // zero out buffer
-        println!("interning with dictionary array: {:#?}", array);
         let dict_array = array.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
         // grab the keys and values array
         let values = dict_array.values();
@@ -123,16 +125,56 @@ impl<K:ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K>
         Ok(())
     }
     fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        Ok(vec![])
+       let columns: Vec<ScalarValue> = match emit_to {
+        EmitTo::All => {
+            self.unique_dict_value_mapping.clear();
+            mem::take(&mut self.seen_elements)
+        },
+        EmitTo::First(n) => {
+            // drain first n elements, keeping the rest
+            let first_n = self.seen_elements.drain(..n).collect();
+            // shift all remaining group indices down by n
+            self.unique_dict_value_mapping.retain(|_, group_idx| {
+                match group_idx.checked_sub(n) {
+                    Some(new_idx) => {
+                        *group_idx = new_idx;
+                        true
+                    }
+                    // this group was in the first n, remove it
+                    None => false,
+                }
+            });
+            first_n
+        }
+    };
+
+    // convert Vec<ScalarValue> into an ArrayRef
+    let array = ScalarValue::iter_to_array(columns.into_iter())?;
+    Ok(vec![array])
+       
+    }
+    fn clear_shrink(&mut self, num_rows: usize) {
+        self.seen_elements.clear();
+        self.seen_elements.shrink_to(num_rows);
+        self.unique_dict_value_mapping.clear();
+        self.unique_dict_value_mapping.shrink_to(num_rows);
+    }
+}
+fn element_size(dt: &DataType) -> usize {
+    match dt{
+        DataType::Utf8 | DataType::LargeUtf8 => 20, // rough estimate for average string size
+        DataType::Binary | DataType::LargeBinary => 20, // rough estimate for average binary size
+        DataType::Boolean => 1,
+        DataType::Int8 | DataType::UInt8 => 1,
+        DataType::Int16 | DataType::UInt16 => 2,
+        DataType::Int32 | DataType::UInt32 => 4,
+        DataType::Int64 | DataType::UInt64 => 8,
+        _ => 0, // default case for unsupported types
     }
-    fn clear_shrink(&mut self, num_rows: usize) {}
 }
 
 #[cfg(test)]
 mod group_values_trait_test {
-    /*
-     cargo test --package datafusion-physical-plan --lib -- aggregates::group_values::single_group_by::dictionary::group_values_trait_test --nocapture
-    */
     use super::*;
     use arrow::array::{DictionaryArray, StringArray, UInt8Array};
     use std::sync::Arc;
@@ -150,47 +192,39 @@ mod group_values_trait_test {
     }
     /*
     cargo test --package datafusion-physical-plan --lib -- aggregates::group_values::single_group_by::dictionary::group_values_trait_test::test_group_values_dictionary --exact --nocapture --include-ignored
-    */
-    #[test]
-    fn test_group_values_dictionary() {
-        run_groupvalue_test_suite().unwrap();
-    }
-
-    fn run_groupvalue_test_suite(
-    ) -> Result<()> {
-        let tests: Vec<(&str,fn(&mut dyn GroupValues))> = vec![
-           ("test_single_group_all_same_values", basic_functionality::test_single_group_all_same_values),
-           ("test_multiple_groups", basic_functionality::test_multiple_groups),
+    
+    fn run_groupvalue_test_suite() -> Result<()> {
+        let tests: Vec<(&str, fn(&mut dyn GroupValues))> = vec![
+            ("test_single_group_all_same_values", basic_functionality::test_single_group_all_same_values),
+            ("test_multiple_groups", basic_functionality::test_multiple_groups),
             ("test_all_different_values", basic_functionality::test_all_different_values),
             ("test_empty_batch", edge_cases::test_empty_batch),
             ("test_single_row", edge_cases::test_single_row),
             ("test_repeated_pattern", edge_cases::test_repeated_pattern),
-            /* 
-            multi_column::test_multiple_columns_passed,
-            consecutive_batches::test_consecutive_batches_then_emit,
-            consecutive_batches::test_three_consecutive_batches_with_partial_emit,
-            state_management::test_size_grows_after_intern,
-            state_management::test_complex_emit_flow_with_multiple_internS,
-            state_management::test_clear_shrink_resets_state,
-            state_management::test_clear_shrink_with_zero,
-            state_management::test_emit_all_clears_state,
-            state_management::test_emit_first_n,
-            state_management::test_complex_emit_flow_with_multiple_internS,
-            data_correctness::test_group_assignment_order,
-            data_correctness::test_groups_vector_correctness_first_appearance,
-            data_correctness::test_groups_vector_sequential_assignment,
-            data_correctness::test_emit_partial_preserves_state,
-            data_correctness::test_emit_restores_intern_ability,
-            */
+            ("test_multiple_columns_passed", multi_column::test_multiple_columns_passed),
+            ("test_consecutive_batches_then_emit", consecutive_batches::test_consecutive_batches_then_emit),
+            ("test_three_consecutive_batches_with_partial_emit", consecutive_batches::test_three_consecutive_batches_with_partial_emit),
+            ("test_size_grows_after_intern", state_management::test_size_grows_after_intern),
+            ("test_complex_emit_flow_with_multiple_internS", state_management::test_complex_emit_flow_with_multiple_internS),
+            ("test_clear_shrink_resets_state", state_management::test_clear_shrink_resets_state),
+            ("test_clear_shrink_with_zero", state_management::test_clear_shrink_with_zero),
+            ("test_emit_all_clears_state", state_management::test_emit_all_clears_state),
+            ("test_emit_first_n", state_management::test_emit_first_n),
+            ("test_group_assignment_order", data_correctness::test_group_assignment_order),
+            ("test_groups_vector_correctness_first_appearance", data_correctness::test_groups_vector_correctness_first_appearance),
+            ("test_groups_vector_sequential_assignment", data_correctness::test_groups_vector_sequential_assignment),
+            ("test_emit_partial_preserves_state", data_correctness::test_emit_partial_preserves_state),
+            ("test_emit_restores_intern_ability", data_correctness::test_emit_restores_intern_ability),
         ];
-        for (name, test_functions) in tests {
+        for (name, test_function) in tests {
             let mut group_values = GroupValuesDictionary::<arrow::datatypes::UInt8Type>::new(&DataType::Utf8);
             println!("Running test: {name}");
-            test_functions(&mut group_values);
+            test_function(&mut group_values);
         }
 
         Ok(())
     }
+    */
 
     mod basic_functionality {
         use super::*;
@@ -364,13 +398,11 @@ mod group_values_trait_test {
             group_values_trait_obj
                 .intern(&[batch2], &mut groups_vector2)
                 .unwrap();
-
             assert_eq!(group_values_trait_obj.len(), 3);
             assert_eq!(groups_vector2.len(), 3);
 
             let result = group_values_trait_obj.emit(EmitTo::All).unwrap();
             assert_eq!(result.len(), 1);
-
             assert!(group_values_trait_obj.is_empty());
         }
 
@@ -397,16 +429,28 @@ mod group_values_trait_test {
                 .unwrap();
             assert_eq!(group_values_trait_obj.len(), 3);
 
-            let batch3 = create_dict_array(vec![2, 3], vec!["c", "d"]);
+            let batch3 = create_dict_array(vec![0, 1,0,1,1,1,1,1,1,0,1,1,0,1,2,1,2], vec!["c", "d","e"]);
             let mut groups_vector3 = Vec::new();
             group_values_trait_obj
                 .intern(&[batch3], &mut groups_vector3)
                 .unwrap();
-            assert_eq!(group_values_trait_obj.len(), 4);
+            assert_eq!(group_values_trait_obj.len(), 5);
 
             let result = group_values_trait_obj.emit(EmitTo::All).unwrap();
             assert_eq!(result.len(), 1);
             assert!(group_values_trait_obj.is_empty());
+            result.iter().for_each(|array| {
+                let string_array = array.as_any().downcast_ref::<StringArray>().unwrap();
+                let values: Vec<String> = (0..string_array.len())
+                    .map(|i| string_array.value(i).to_string())
+                    .collect();
+                let unexpected_values: Vec<&String> = values.iter().filter(|v| **v != "a" && **v != "b" && **v != "c" && **v != "d" && **v != "e").collect();
+                assert!(
+                    unexpected_values.is_empty(),
+                    "Emitted unexpected values: {:#?}",
+                    unexpected_values
+                );
+            });
         }
 
         #[test]
@@ -422,7 +466,6 @@ mod group_values_trait_test {
         fn test_initial_state_is_empty(group_values_trait_obj: &dyn GroupValues) {
             assert!(group_values_trait_obj.is_empty());
             assert_eq!(group_values_trait_obj.len(), 0);
-            assert_eq!(group_values_trait_obj.size(), 0);
         }
 
         #[test]
@@ -756,7 +799,7 @@ mod group_values_trait_test {
 
         #[test]
         fn run_test_emit_partial_preserves_state() {
-            let mut group_values = GroupValuesDictionary::<arrow::datatypes::Int8Type>::new(&DataType::Utf8);
+            let mut group_values = GroupValuesDictionary::<arrow::datatypes::UInt8Type>::new(&DataType::Utf8);
             test_emit_partial_preserves_state(&mut group_values);
         }