perf: reuse cached dictionary value hashes in vectorized_append

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

@@ -54,12 +54,18 @@ pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync>
     null_inner_slot: Option<usize>,
     /// Hash seed — must match `create_hashes` so hashes are consistent across calls.
     random_state: RandomState,
-    /// Reusable scratch buffer mapping `val_idx → inner_slot` across batches.
+    /// Maps `val_idx → inner_slot` for the values array in `cached_values`.
+    ///
+    /// Valid only while `cached_values` is `Some` and matches the incoming
+    /// values array; `usize::MAX` means "not yet resolved". Entries are only
+    /// ever filled in, never invalidated, because `inner` slots are stable
+    /// under append. `take_n` remaps `inner` and therefore drops the cache.
     val_to_inner: Vec<usize>,
-    /// Reusable hash buffer for the dictionary values array.
+    /// Hashes of the values array in `cached_values`, one per `val_idx`.
     val_hashes: Vec<u64>,
-    /// The last `dict.values()` Arc hashed in `append_val`. When the incoming
-    /// values array is `ptr_eq` to this, `val_hashes` can be reused directly.
+    /// The values array that `val_hashes` and `val_to_inner` were built for.
+    /// When an incoming `dict.values()` is `Arc::ptr_eq` to this, both caches
+    /// are reused instead of being rebuilt. See [`Self::sync_value_cache`].
     cached_values: Option<ArrayRef>,
     _phantom: PhantomData<K>,
 }
@@ -175,6 +181,28 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
         .unwrap();
     }
 
+    /// Makes `val_hashes` and `val_to_inner` correspond to `dict_values`,
+    /// reusing them when they were already built for this same values array.
+    ///
+    /// `take`, `filter` and repartitioning clone the values `Arc` and rewrite
+    /// only the keys, so batches in a partition usually share one values
+    /// array; without this the per-batch cost is O(dictionary cardinality)
+    /// rather than O(rows). `take_n` remaps `inner` and clears
+    /// `cached_values`, forcing a rebuild.
+    fn sync_value_cache(&mut self, dict_values: &ArrayRef) {
+        if self
+            .cached_values
+            .as_ref()
+            .is_some_and(|c| Arc::ptr_eq(c, dict_values))
+        {
+            return;
+        }
+        self.hash_values(dict_values);
+        self.val_to_inner.clear();
+        self.val_to_inner.resize(dict_values.len(), usize::MAX);
+        self.cached_values = Some(Arc::clone(dict_values));
+    }
+
     fn find_or_insert_value(
         &mut self,
         dict_values: &ArrayRef,
@@ -302,24 +330,7 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
             }
             Some(val_idx) => {
                 let dict_values = dict.values();
-                // check if the dictionary values array we are hashing was already seen.
-                // if its arc was already stored we dont need to rehash the entire array again
-                // if its new hash the entire array and store an arc ptr for future use
-                let cache_hit = self
-                    .cached_values
-                    .as_ref()
-                    .is_some_and(|c| Arc::ptr_eq(c, dict_values));
-                if !cache_hit {
-                    self.val_hashes.clear();
-                    self.val_hashes.resize(dict_values.len(), 0);
-                    create_hashes(
-                        std::slice::from_ref(dict_values),
-                        &self.random_state,
-                        &mut self.val_hashes,
-                    )
-                    .unwrap();
-                    self.cached_values = Some(Arc::clone(dict_values));
-                }
+                self.sync_value_cache(dict_values);
                 self.find_or_insert_value(dict_values, val_idx, self.val_hashes[val_idx])?
             }
         };
@@ -401,11 +412,12 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
         let dict = array.as_dictionary::<K>();
         let dict_keys = dict.keys();
         let dict_values = dict.values();
-        let num_distinct = dict_values.len();
 
-        self.hash_values(dict_values);
-        self.val_to_inner.clear();
-        self.val_to_inner.resize(num_distinct, usize::MAX);
+        // Reuses `val_hashes`/`val_to_inner` when this batch carries the same
+        // values array as the last one, which is the common case downstream of
+        // a repartition or filter. Previously this unconditionally re-hashed
+        // and re-initialised both buffers, costing O(D) per batch.
+        self.sync_value_cache(dict_values);
 
         self.group_to_inner.try_reserve(rows.len()).map_err(|e| {
             DataFusionError::ArrowError(
@@ -918,4 +930,134 @@ mod tests {
         assert_eq!(str_values(&out)[127], Some("v127".into()));
         assert_eq!(out.as_dictionary::<Int8Type>().values().len(), 128);
     }
+
+    fn dict_with_values(keys: &[Option<i32>], values: &ArrayRef) -> ArrayRef {
+        Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(keys.to_vec()),
+            Arc::clone(values),
+        ))
+    }
+
+    /// A cache rebuild resets every `val_to_inner` entry to `usize::MAX`, so
+    /// an entry resolved by an earlier batch surviving a later one proves the
+    /// dictionary was not re-hashed in between.
+    #[test]
+    fn cache_survives_batches_sharing_a_values_array() {
+        let values: ArrayRef =
+            Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")]));
+        let mut col = utf8_col();
+
+        col.vectorized_append(&dict_with_values(&[Some(0)], &values), &[0])
+            .unwrap();
+        let slot_a = col.val_to_inner[0];
+        assert_ne!(slot_a, usize::MAX);
+
+        col.vectorized_append(&dict_with_values(&[Some(1)], &values), &[0])
+            .unwrap();
+
+        assert_eq!(col.val_to_inner[0], slot_a, "cache was rebuilt");
+        assert_ne!(col.val_to_inner[1], usize::MAX);
+        assert_eq!(
+            col.val_to_inner[2],
+            usize::MAX,
+            "unreferenced value resolved"
+        );
+        assert!(
+            col.cached_values
+                .as_ref()
+                .is_some_and(|c| Arc::ptr_eq(c, &values))
+        );
+        assert_eq!(col.inner.len(), 2);
+    }
+
+    /// `take` clones the values `Arc` and rewrites only the keys, which is the
+    /// property the `ptr_eq` check relies on. Pin that arrow behaviour.
+    #[test]
+    fn take_preserves_values_arc_so_cache_hits() {
+        let values: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")]));
+        let batch = dict_with_values(&[Some(0), Some(1), Some(0)], &values);
+        let taken = take(&*batch, &Int32Array::from(vec![2, 0]), None).unwrap();
+        assert!(
+            Arc::ptr_eq(taken.as_dictionary::<Int32Type>().values(), &values),
+            "take must not copy the values array"
+        );
+
+        let mut col = utf8_col();
+        col.vectorized_append(&batch, &[0, 1, 2]).unwrap();
+        let slot_b = col.val_to_inner[1];
+        col.vectorized_append(&taken, &[0, 1]).unwrap();
+
+        assert_eq!(col.val_to_inner[1], slot_b, "cache was rebuilt");
+        assert_eq!(col.inner.len(), 2);
+    }
+
+    /// A genuinely different values array must rebuild the cache, and equal
+    /// values across the two arrays must still dedup to one inner slot.
+    #[test]
+    fn distinct_values_arcs_rebuild_cache_but_still_dedup() {
+        let mut col = utf8_col();
+        let first = i32_dict(&[Some(0)], &[Some("a")]);
+        let second = i32_dict(&[Some(0)], &[Some("a")]);
+        col.vectorized_append(&first, &[0]).unwrap();
+        col.vectorized_append(&second, &[0]).unwrap();
+
+        assert!(
+            col.cached_values.as_ref().is_some_and(|c| Arc::ptr_eq(
+                c,
+                second.as_dictionary::<Int32Type>().values()
+            )),
+            "cache should track the most recent values array"
+        );
+        assert_eq!(col.inner.len(), 1, "'a' must dedup across both arrays");
+        assert_eq!(col.len(), 2);
+    }
+
+    /// `take_n` remaps inner slots, so the cached `val_idx → inner_slot` map
+    /// must not survive it.
+    #[test]
+    fn take_n_invalidates_value_cache() {
+        let values: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")]));
+        let mut col = utf8_col();
+        col.vectorized_append(&dict_with_values(&[Some(0), Some(1)], &values), &[0, 1])
+            .unwrap();
+
+        let _emitted = col.take_n(1);
+        assert!(
+            col.cached_values.is_none(),
+            "take_n remaps inner slots and must drop the cache"
+        );
+
+        col.vectorized_append(&dict_with_values(&[Some(0), Some(1)], &values), &[0, 1])
+            .unwrap();
+
+        let out = Box::new(col).build();
+        let d = out.as_dictionary::<Int32Type>();
+        let vals = d.values().as_string::<i32>();
+        let got: Vec<&str> = d
+            .keys()
+            .iter()
+            .map(|k| vals.value(k.unwrap() as usize))
+            .collect();
+        assert_eq!(got, vec!["b", "a", "b"]);
+    }
+
+    /// `append_val` fills `val_hashes` but not `val_to_inner`; mixing it with
+    /// `vectorized_append` must not duplicate a value or reuse a stale slot.
+    #[test]
+    fn scalar_and_vectorized_append_share_cache_consistently() {
+        let values: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), Some("b")]));
+        let batch = dict_with_values(&[Some(0), Some(1)], &values);
+        let mut col = utf8_col();
+
+        col.append_val(&batch, 0).unwrap();
+        assert!(
+            col.cached_values
+                .as_ref()
+                .is_some_and(|c| Arc::ptr_eq(c, &values))
+        );
+
+        col.vectorized_append(&batch, &[0, 1]).unwrap();
+        assert_eq!(col.inner.len(), 2, "'a' must not be duplicated");
+        assert_eq!(col.len(), 3);
+    }
 }