updated normalization fn to use a hashmap cache instead of doing a n^2 loop

Author: Rich-T-kid

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

@@ -203,7 +203,8 @@ impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
             ))),
         }
     }
-    fn normalize_dict_array(
+    #[inline]
+    fn normalize_dict_linear(
         values: &ArrayRef,
         key_array: &PrimitiveArray<K>,
     ) -> (ArrayRef, Vec<Option<usize>>) {
@@ -249,6 +250,61 @@ impl<K: ArrowDictionaryKeyType + Send> GroupValuesDictionary<K> {
 
         (new_values, new_keys)
     }
+    #[inline]
+    pub fn normalize_dict_hash(
+        values: &ArrayRef,
+        key_array: &PrimitiveArray<K>,
+    ) -> (ArrayRef, Vec<Option<usize>>) {
+        let mut key_to_canonical: Vec<Option<usize>> = vec![None; values.len()];
+        let mut bytes_to_canonical: HashMap<Vec<u8>, usize> = HashMap::new();
+        let mut canonical_indices: Vec<usize> = Vec::new();
+
+        let new_keys: Vec<Option<usize>> = (0..key_array.len())
+            .map(|i| {
+                if key_array.is_null(i) {
+                    return None;
+                }
+                let old_key = key_array.value(i).to_usize().unwrap();
+                if let Some(canonical) = key_to_canonical[old_key] {
+                    return Some(canonical);
+                }
+                let raw = Self::get_raw_bytes(values, old_key).to_vec();
+                let canonical = match bytes_to_canonical.get(&raw) {
+                    Some(&idx) => idx,
+                    None => {
+                        let idx = canonical_indices.len();
+                        canonical_indices.push(old_key);
+                        bytes_to_canonical.insert(raw, idx);
+                        idx
+                    }
+                };
+                key_to_canonical[old_key] = Some(canonical);
+                Some(canonical)
+            })
+            .collect();
+
+        let indices = UInt64Array::from(
+            canonical_indices
+                .iter()
+                .map(|&i| i as u64)
+                .collect::<Vec<_>>(),
+        );
+        let new_values = arrow::compute::take(values.as_ref(), &indices, None).unwrap();
+        (new_values, new_keys)
+    }
+    #[inline]
+    fn normalize_dict_array(
+        values: &ArrayRef,
+        key_array: &PrimitiveArray<K>,
+    ) -> (ArrayRef, Vec<Option<usize>>) {
+        if values.len() < 20 {
+            //benchmarks show that for small arrays a nested loops works better due to cache locality and lower overhead,
+            Self::normalize_dict_linear(values, key_array)
+        } else {
+            //for larger arrays the hash-based approach is much faster due to avoiding redundant comparisons
+            Self::normalize_dict_hash(values, key_array)
+        }
+    }
 }
 
 impl<K: ArrowDictionaryKeyType + Send> GroupValues for GroupValuesDictionary<K> {