add cache to concat pointers to avoid un-needed allocations

Author: Rich-T-kid

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

@@ -25,10 +25,15 @@ use datafusion_common::Result;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
+/// A [`GroupColumn`] implementation for dictionary arrays.
+/// wraps an inner [`GroupColumn`] that stores the dictionary values, and uses a null sentinel for null keys.
 pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> {
     inner: Box<dyn GroupColumn>,
+    // unary array used as a null sentinel for dictionary keys that are null
     null_array: ArrayRef,
     _phantom: PhantomData<K>,
+    values_ptr: Option<ArrayRef>,
+    combined_ptr: Option<ArrayRef>,
 }
 
 impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
@@ -38,6 +43,8 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
             inner,
             null_array,
             _phantom: PhantomData,
+            values_ptr: None,
+            combined_ptr: None,
         }
     }
 
@@ -115,9 +122,19 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
                 equal_to_results,
             );
         } else {
-            // Null keys must be routed to a null sentinel.
-            let combined = concat(&[dict_values.as_ref(), self.null_array.as_ref()])
-                .expect("concat of dict values and null sentinel should not fail");
+            // Null keys must be routed to a null sentinel. A cheap ptr check
+            // avoids an O(values) concat and allocation when the dictionary
+            // array is the same across calls (populated by vectorized_append).
+            let combined = {
+                if let Some(values_ptr) = self.values_ptr.as_ref()
+                    && Arc::ptr_eq(values_ptr, dict_values)
+                {
+                    self.combined_ptr.clone().unwrap()
+                } else {
+                    concat(&[dict_values.as_ref(), self.null_array.as_ref()])
+                        .expect("concat of dict values and null sentinel should not fail")
+                }
+            };
             let null_sentinel_index = combined.len() - 1;
             let value_indices: Vec<usize> = rhs_rows
                 .iter()
@@ -143,9 +160,21 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
                 .collect();
             self.inner.vectorized_append(dict.values(), &value_indices)
         } else {
-            let combined_with_null_sentinel =
-                concat(&[dict.values(), self.null_array.as_ref()])
-                    .expect("concat of dict values and null sentinel should not fail");
+            let combined_with_null_sentinel = {
+                if let Some(values_ptr) = self.values_ptr.as_ref()
+                    && Arc::ptr_eq(values_ptr, dict.values())
+                {
+                    self.combined_ptr.clone().unwrap()
+                } else {
+                    let combined = concat(&[dict.values(), self.null_array.as_ref()])
+                        .expect(
+                            "concat of dict values and null sentinel should not fail",
+                        );
+                    self.values_ptr = Some(Arc::clone(dict.values()));
+                    self.combined_ptr = Some(Arc::clone(&combined));
+                    combined
+                }
+            };
             let null_sentinel_index = combined_with_null_sentinel.len() - 1;
             let value_indices: Vec<usize> = rows
                 .iter()