WIP: Reuse dictionary value hashes in hash joins

datafusion/common/src/hash_utils.rs

@@ -17,6 +17,10 @@ use foldhash::fast::FixedState;
 #[cfg(not(feature = "force_hash_collisions"))]
 use itertools::Itertools;
 use std::hash::{BuildHasher, Hash, Hasher};
+use std::mem::size_of;
+#[cfg(not(feature = "force_hash_collisions"))]
+use std::sync::Arc;
+use std::sync::Weak;
 
 /// [`RandomState`] is optimized for speed and suitable for hash tables and
 /// bloom filters. [`QualityRandomState`] is optimized for statistical quality
@@ -28,6 +32,134 @@ use std::hash::{BuildHasher, Hash, Hasher};
 pub type RandomState = FixedState;
 pub type QualityRandomState = foldhash::quality::FixedState;
 
+type DictionaryHashEntry = (Weak<dyn Array>, Vec<u64>);
+
+/// Hashes arrays, reusing dictionary value hashes by column when possible.
+#[derive(Debug)]
+pub struct CachingArrayHasher {
+    /// Stored with the cache so value hashes cannot cross hash seeds.
+    random_state: RandomState,
+    /// One entry per column. Weak identity avoids retaining dictionary values.
+    dictionaries: Vec<Option<DictionaryHashEntry>>,
+}
+
+impl CachingArrayHasher {
+    /// Creates one cache slot per column in `schema`.
+    ///
+    /// For each dictionary column, the cache tracks the values array with a
+    /// weak pointer, so it does not keep the values alive, and owns a reusable
+    /// value-hash buffer. [`Self::create_hashes`] reuses those hashes while the
+    /// values `Arc` is unchanged, recomputes them when it changes, and scatters
+    /// them through the dictionary keys into the row hashes. Non-dictionary
+    /// columns use the regular array hashing path.
+    pub fn new(schema: &Schema, random_state: RandomState) -> Self {
+        Self {
+            random_state,
+            dictionaries: schema.fields().iter().map(|_| None).collect(),
+        }
+    }
+
+    /// Hashes `arrays`, reusing dictionary value hashes while each column's
+    /// values array is unchanged.
+    pub fn create_hashes<'a, I, T>(
+        &mut self,
+        arrays: I,
+        hashes_buffer: &'a mut [u64],
+    ) -> Result<&'a mut [u64]>
+    where
+        I: IntoIterator<Item = T>,
+        I::IntoIter: ExactSizeIterator,
+        T: AsDynArray,
+    {
+        let arrays = arrays.into_iter();
+        if arrays.len() != self.dictionaries.len() {
+            return _internal_err!(
+                "CachingArrayHasher expected {} arrays, got {}",
+                self.dictionaries.len(),
+                arrays.len()
+            );
+        }
+        for (index, array) in arrays.enumerate() {
+            self.hash_array(index, array.as_dyn_array(), hashes_buffer, index != 0)?;
+        }
+        Ok(hashes_buffer)
+    }
+
+    fn hash_array(
+        &mut self,
+        index: usize,
+        array: &dyn Array,
+        hashes_buffer: &mut [u64],
+        rehash: bool,
+    ) -> Result<()> {
+        #[cfg(feature = "force_hash_collisions")]
+        {
+            let _ = index;
+            hash_single_array(array, &self.random_state, hashes_buffer, rehash)?;
+        }
+
+        #[cfg(not(feature = "force_hash_collisions"))]
+        {
+            let dictionary = &mut self.dictionaries[index];
+            if !matches!(array.data_type(), DataType::Dictionary(_, _)) {
+                *dictionary = None;
+                hash_single_array(array, &self.random_state, hashes_buffer, rehash)?;
+                return Ok(());
+            }
+
+            downcast_dictionary_array! {
+                array => {
+                    let values = array.values();
+                    if !dictionary.as_ref().is_some_and(|(cached, _)| {
+                        std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values))
+                    }) {
+                        let mut hashes = dictionary
+                            .take()
+                            .map(|(_, hashes)| hashes)
+                            .unwrap_or_default();
+                        hashes.resize(values.len(), 0);
+                        // Fresh hashing leaves null value slots unchanged.
+                        hashes.fill(0);
+                        hash_single_array(
+                            values.as_ref(),
+                            &self.random_state,
+                            &mut hashes,
+                            false,
+                        )?;
+                        *dictionary = Some((Arc::downgrade(values), hashes));
+                    }
+                    dispatch_dictionary_scatter(
+                        array,
+                        &dictionary.as_ref().unwrap().1,
+                        hashes_buffer,
+                        rehash,
+                    );
+                }
+                _ => unreachable!(),
+            };
+        }
+
+        Ok(())
+    }
+
+    /// Returns the current dictionary values and hashes for `index`, if any.
+    pub fn dictionary_hashes(&self, index: usize) -> Option<(ArrayRef, &[u64])> {
+        let (values, hashes) = self.dictionaries.get(index)?.as_ref()?;
+        Some((values.upgrade()?, hashes))
+    }
+
+    /// Returns the allocated size of the cache.
+    pub fn size(&self) -> usize {
+        self.dictionaries.capacity() * size_of::<Option<DictionaryHashEntry>>()
+            + self
+                .dictionaries
+                .iter()
+                .flatten()
+                .map(|entry| entry.1.capacity() * size_of::<u64>())
+                .sum::<usize>()
+    }
+}
+
 /// Fixed quality hash state used by HyperLogLog sketches.
 ///
 /// The seed is part of the HLL wire/storage semantics: serialized sketches only
@@ -1357,6 +1489,194 @@ mod tests {
         }
     }
 
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn string_dictionary(keys: Vec<i8>, values: &ArrayRef) -> ArrayRef {
+        Arc::new(DictionaryArray::<Int8Type>::new(
+            Int8Array::from(keys),
+            Arc::clone(values),
+        ))
+    }
+
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn caching_hasher(
+        columns: &[ArrayRef],
+        random_state: &RandomState,
+    ) -> CachingArrayHasher {
+        let fields = columns
+            .iter()
+            .enumerate()
+            .map(|(index, column)| {
+                Field::new(format!("column_{index}"), column.data_type().clone(), true)
+            })
+            .collect::<Vec<_>>();
+        CachingArrayHasher::new(&Schema::new(fields), random_state.clone())
+    }
+
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn assert_hashes_match(
+        hasher: &mut CachingArrayHasher,
+        columns: &[ArrayRef],
+        random_state: &RandomState,
+    ) -> Result<()> {
+        let mut expected = vec![0; columns[0].len()];
+        create_hashes(columns, random_state, &mut expected)?;
+
+        let mut actual = vec![0; columns[0].len()];
+        hasher.create_hashes(columns, &mut actual)?;
+        assert_eq!(actual, expected);
+        Ok(())
+    }
+
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn assert_cached_values(
+        hasher: &CachingArrayHasher,
+        column_index: usize,
+        expected_values: &ArrayRef,
+    ) -> *const u64 {
+        let (cached_values, value_hashes) =
+            hasher.dictionary_hashes(column_index).unwrap();
+        assert!(Arc::ptr_eq(&cached_values, expected_values));
+        value_hashes.as_ptr()
+    }
+
+    #[cfg(not(feature = "force_hash_collisions"))]
+    #[test]
+    fn caching_array_hasher_hashes_regular_and_dictionary_columns() -> Result<()> {
+        let random_state = RandomState::with_seed(0);
+        let dictionary_values: ArrayRef =
+            Arc::new(StringArray::from(vec!["active", "inactive"]));
+        let first_batch = vec![
+            Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef,
+            string_dictionary(vec![0, 1, 0], &dictionary_values),
+        ];
+        let mut hasher = caching_hasher(&first_batch, &random_state);
+        assert_hashes_match(&mut hasher, &first_batch, &random_state)?;
+
+        assert!(hasher.dictionary_hashes(0).is_none());
+        let value_hash_allocation = assert_cached_values(&hasher, 1, &dictionary_values);
+
+        // A later batch has new plain values and dictionary keys, but shares
+        // the dictionary values array and therefore its cached hashes.
+        let second_batch = vec![
+            Arc::new(Int32Array::from(vec![40, 50])) as ArrayRef,
+            string_dictionary(vec![1, 0], &dictionary_values),
+        ];
+        assert_hashes_match(&mut hasher, &second_batch, &random_state)?;
+        assert_eq!(
+            assert_cached_values(&hasher, 1, &dictionary_values),
+            value_hash_allocation
+        );
+        Ok(())
+    }
+
+    #[cfg(not(feature = "force_hash_collisions"))]
+    #[test]
+    fn caching_array_hasher_updates_dictionary_columns_independently() -> Result<()> {
+        let random_state = RandomState::with_seed(0);
+        let country_values: ArrayRef = Arc::new(StringArray::from(vec!["us", "uk"]));
+        let status_values: ArrayRef =
+            Arc::new(StringArray::from(vec!["active", "inactive"]));
+        let first_batch = vec![
+            string_dictionary(vec![0, 1, 0], &country_values),
+            string_dictionary(vec![1, 0, 1], &status_values),
+        ];
+        let mut hasher = caching_hasher(&first_batch, &random_state);
+        assert_hashes_match(&mut hasher, &first_batch, &random_state)?;
+        let country_hash_allocation = assert_cached_values(&hasher, 0, &country_values);
+        let status_hash_allocation = assert_cached_values(&hasher, 1, &status_values);
+
+        // Aggregation batches can rebuild one dictionary while another group
+        // column continues to share its values array.
+        let new_country_values: ArrayRef =
+            Arc::new(StringArray::from(vec![Some("ca"), None]));
+        let second_batch = vec![
+            string_dictionary(vec![1, 0], &new_country_values),
+            string_dictionary(vec![0, 1], &status_values),
+        ];
+        assert_hashes_match(&mut hasher, &second_batch, &random_state)?;
+        assert_eq!(
+            assert_cached_values(&hasher, 0, &new_country_values),
+            country_hash_allocation
+        );
+        assert_eq!(
+            assert_cached_values(&hasher, 1, &status_values),
+            status_hash_allocation
+        );
+        assert_eq!(hasher.dictionary_hashes(0).unwrap().1[1], 0);
+
+        // Reverse the transition: the first dictionary stays shared while the
+        // second receives a new values array.
+        let new_status_values: ArrayRef =
+            Arc::new(StringArray::from(vec!["pending", "complete"]));
+        let third_batch = vec![
+            string_dictionary(vec![0, 1], &new_country_values),
+            string_dictionary(vec![1, 0], &new_status_values),
+        ];
+        assert_hashes_match(&mut hasher, &third_batch, &random_state)?;
+        assert_eq!(
+            assert_cached_values(&hasher, 0, &new_country_values),
+            country_hash_allocation
+        );
+        assert_eq!(
+            assert_cached_values(&hasher, 1, &new_status_values),
+            status_hash_allocation
+        );
+        Ok(())
+    }
+
+    #[cfg(not(feature = "force_hash_collisions"))]
+    #[test]
+    fn caching_array_hasher_tracks_repeated_changes_with_regular_column() -> Result<()> {
+        let random_state = RandomState::with_seed(0);
+        let region_values_v1: ArrayRef =
+            Arc::new(StringArray::from(vec!["east", "west"]));
+        let category_values_v1: ArrayRef =
+            Arc::new(StringArray::from(vec!["hardware", "software"]));
+        let first_batch = vec![
+            string_dictionary(vec![0, 1, 0], &region_values_v1),
+            string_dictionary(vec![1, 0, 1], &category_values_v1),
+            Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef,
+        ];
+        let mut hasher = caching_hasher(&first_batch, &random_state);
+        assert_hashes_match(&mut hasher, &first_batch, &random_state)?;
+
+        // This mirrors repeated aggregation input: dictionary group columns
+        // can change independently while regular group columns change normally.
+        let region_values_v2: ArrayRef =
+            Arc::new(StringArray::from(vec!["north", "south"]));
+        let second_batch = vec![
+            string_dictionary(vec![1, 0, 1], &region_values_v2),
+            string_dictionary(vec![0, 1, 0], &category_values_v1),
+            Arc::new(Int32Array::from(vec![40, 50, 60])) as ArrayRef,
+        ];
+        assert_hashes_match(&mut hasher, &second_batch, &random_state)?;
+
+        let category_values_v2: ArrayRef =
+            Arc::new(StringArray::from(vec!["services", "support"]));
+        let third_batch = vec![
+            string_dictionary(vec![0, 1], &region_values_v2),
+            string_dictionary(vec![1, 0], &category_values_v2),
+            Arc::new(Int32Array::from(vec![70, 80])) as ArrayRef,
+        ];
+        assert_hashes_match(&mut hasher, &third_batch, &random_state)?;
+
+        // Replacing both dictionaries together must also update both cache
+        // slots without affecting the regular column.
+        let region_values_v3: ArrayRef =
+            Arc::new(StringArray::from(vec!["central", "coastal", "remote"]));
+        let category_values_v3: ArrayRef = Arc::new(StringArray::from(vec!["other"]));
+        let fourth_batch = vec![
+            string_dictionary(vec![2, 0, 1], &region_values_v3),
+            string_dictionary(vec![0, 0, 0], &category_values_v3),
+            Arc::new(Int32Array::from(vec![90, 100, 110])) as ArrayRef,
+        ];
+        assert_hashes_match(&mut hasher, &fourth_batch, &random_state)?;
+        assert_cached_values(&hasher, 0, &region_values_v3);
+        assert_cached_values(&hasher, 1, &category_values_v3);
+        assert!(hasher.dictionary_hashes(2).is_none());
+        Ok(())
+    }
+
     #[test]
     fn create_hashes_for_decimal_array() -> Result<()> {
         let array = vec![1, 2, 3, 4]