Skip to content

Commit 81d2bb6

Browse files
Rich-T-kidRich-T-kid
committed
Avoid rehashing shared dictionaries in RepartitionExec
Cache dictionary value hashes per input partition while batches share the same values array. Account retained hashes with the memory pool and fall back to uncached hashing if the reservation is unavailable.
1 parent acbad14 commit 81d2bb6

2 files changed

Lines changed: 407 additions & 15 deletions

File tree

‎datafusion/common/src/hash_utils.rs‎

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ use foldhash::fast::FixedState;
3030
#[cfg(not(feature = "force_hash_collisions"))]
3131
use itertools::Itertools;
3232
use std::hash::{BuildHasher, Hash, Hasher};
33+
#[cfg(not(feature = "force_hash_collisions"))]
34+
use std::sync::{Arc, Weak};
3335

3436
/// [`RandomState`] is optimized for speed and suitable for hash tables and
3537
/// bloom filters. [`QualityRandomState`] is optimized for statistical quality
@@ -679,6 +681,174 @@ fn hash_dictionary<K: ArrowDictionaryKeyType>(
679681
Ok(())
680682
}
681683

684+
#[cfg(not(feature = "force_hash_collisions"))]
685+
#[derive(Default)]
686+
struct DictionaryHashCacheEntry {
687+
/// Weak reference avoids retaining an input dictionary after its batch is released.
688+
values: Option<Weak<dyn Array>>,
689+
/// Hash for each value in `values`.
690+
hashes: Vec<u64>,
691+
/// Whether `hashes` has been populated for `values`.
692+
ready: bool,
693+
/// Number of cache misses, used to verify reuse.
694+
#[cfg(test)]
695+
rebuilds: usize,
696+
}
697+
698+
#[cfg(not(feature = "force_hash_collisions"))]
699+
impl DictionaryHashCacheEntry {
700+
fn get_or_insert<'a>(
701+
&'a mut self,
702+
values: &ArrayRef,
703+
random_state: &impl HashState,
704+
) -> Result<Option<&'a [u64]>> {
705+
let is_cached = self.values.as_ref().is_some_and(|cached| {
706+
std::ptr::addr_eq(cached.as_ptr(), Arc::as_ptr(values))
707+
});
708+
if !is_cached {
709+
self.values = Some(Arc::downgrade(values));
710+
self.hashes = Vec::new();
711+
self.ready = false;
712+
return Ok(None);
713+
}
714+
715+
if !self.ready {
716+
self.hashes.clear();
717+
self.hashes.resize(values.len(), 0);
718+
create_hashes([values], random_state, &mut self.hashes)?;
719+
self.ready = true;
720+
#[cfg(test)]
721+
{
722+
self.rebuilds += 1;
723+
}
724+
}
725+
Ok(Some(&self.hashes))
726+
}
727+
728+
fn size(&self) -> usize {
729+
self.hashes.capacity() * size_of::<u64>()
730+
}
731+
}
732+
733+
/// Caches dictionary value hashes by column and values array identity.
734+
///
735+
/// Arrow arrays are immutable, so matching [`Arc`] identity is sufficient to
736+
/// reuse the cached hashes.
737+
///
738+
/// This is an internal API used by physical operators that hash multiple batches
739+
/// sharing the same dictionary values.
740+
#[doc(hidden)]
741+
pub struct DictionaryHashCache {
742+
/// Hash state used to build all cached hashes.
743+
random_state: RandomState,
744+
/// One cache entry per hashed column.
745+
#[cfg(not(feature = "force_hash_collisions"))]
746+
entries: Vec<DictionaryHashCacheEntry>,
747+
}
748+
749+
impl DictionaryHashCache {
750+
/// Creates a cache for hashes produced by `random_state`.
751+
pub fn new(random_state: RandomState) -> Self {
752+
Self {
753+
random_state,
754+
#[cfg(not(feature = "force_hash_collisions"))]
755+
entries: Vec::new(),
756+
}
757+
}
758+
759+
/// Hashes `arrays`, reusing dictionary value hashes while their values
760+
/// [`Arc`] remains unchanged.
761+
pub fn create_hashes<'a>(
762+
&mut self,
763+
arrays: &[ArrayRef],
764+
hashes_buffer: &'a mut [u64],
765+
) -> Result<&'a mut [u64]> {
766+
#[cfg(feature = "force_hash_collisions")]
767+
return create_hashes(arrays, &self.random_state, hashes_buffer);
768+
769+
#[cfg(not(feature = "force_hash_collisions"))]
770+
{
771+
if !arrays
772+
.iter()
773+
.any(|array| matches!(array.data_type(), DataType::Dictionary(_, _)))
774+
{
775+
self.entries = Vec::new();
776+
return create_hashes(arrays, &self.random_state, hashes_buffer);
777+
}
778+
779+
self.entries
780+
.resize_with(arrays.len(), DictionaryHashCacheEntry::default);
781+
782+
for (index, array) in arrays.iter().enumerate() {
783+
let rehash = index != 0;
784+
if matches!(array.data_type(), DataType::Dictionary(_, _)) {
785+
let entry = &mut self.entries[index];
786+
let array = array.as_ref();
787+
downcast_dictionary_array! {
788+
array => {
789+
if let Some(hashes) = entry.get_or_insert(
790+
array.values(),
791+
&self.random_state,
792+
)? {
793+
dispatch_dictionary_scatter(
794+
array,
795+
hashes,
796+
hashes_buffer,
797+
rehash,
798+
);
799+
} else {
800+
hash_dictionary(
801+
array,
802+
&self.random_state,
803+
hashes_buffer,
804+
rehash,
805+
)?;
806+
}
807+
},
808+
_ => unreachable!()
809+
}
810+
} else {
811+
self.entries[index] = DictionaryHashCacheEntry::default();
812+
hash_single_array(
813+
array.as_ref(),
814+
&self.random_state,
815+
hashes_buffer,
816+
rehash,
817+
)?;
818+
}
819+
}
820+
Ok(hashes_buffer)
821+
}
822+
}
823+
824+
/// Returns the memory retained by cached dictionary value hashes.
825+
pub fn size(&self) -> usize {
826+
#[cfg(feature = "force_hash_collisions")]
827+
return 0;
828+
829+
#[cfg(not(feature = "force_hash_collisions"))]
830+
{
831+
if !self.entries.iter().any(|entry| entry.ready) {
832+
return 0;
833+
}
834+
self.entries.capacity() * size_of::<DictionaryHashCacheEntry>()
835+
+ self
836+
.entries
837+
.iter()
838+
.map(DictionaryHashCacheEntry::size)
839+
.sum::<usize>()
840+
}
841+
}
842+
843+
/// Drops all cached hashes.
844+
pub fn clear(&mut self) {
845+
#[cfg(not(feature = "force_hash_collisions"))]
846+
{
847+
self.entries = Vec::new();
848+
}
849+
}
850+
}
851+
682852
#[cfg(not(feature = "force_hash_collisions"))]
683853
fn hash_dictionary_with_child_hashing<K: ArrowDictionaryKeyType>(
684854
array: &DictionaryArray<K>,
@@ -1353,6 +1523,78 @@ mod tests {
13531523

13541524
use super::*;
13551525

1526+
#[cfg(not(feature = "force_hash_collisions"))]
1527+
fn dictionary_array(keys: Vec<Option<i32>>, values: ArrayRef) -> ArrayRef {
1528+
Arc::new(
1529+
DictionaryArray::<Int32Type>::try_new(Int32Array::from(keys), values)
1530+
.unwrap(),
1531+
)
1532+
}
1533+
1534+
#[cfg(not(feature = "force_hash_collisions"))]
1535+
#[test]
1536+
fn dictionary_hash_cache_reuses_and_invalidates_value_hashes() -> Result<()> {
1537+
let values: ArrayRef =
1538+
Arc::new(StringArray::from(vec![Some("a"), None, Some("ccc")]));
1539+
let first =
1540+
dictionary_array(vec![Some(0), Some(2), None, Some(1)], Arc::clone(&values));
1541+
let second =
1542+
dictionary_array(vec![Some(2), Some(0), Some(2), None], Arc::clone(&values));
1543+
let values_strong_count = Arc::strong_count(&values);
1544+
let random_state = RandomState::with_seed(42);
1545+
let mut cache = DictionaryHashCache::new(random_state.clone());
1546+
1547+
for array in [&first, &second] {
1548+
let mut expected = vec![0; array.len()];
1549+
create_hashes([array], &random_state, &mut expected)?;
1550+
let mut actual = vec![0; array.len()];
1551+
cache.create_hashes(std::slice::from_ref(array), &mut actual)?;
1552+
assert_eq!(actual, expected);
1553+
}
1554+
assert_eq!(cache.entries[0].rebuilds, 1);
1555+
assert_eq!(Arc::strong_count(&values), values_strong_count);
1556+
1557+
let replacement_values: ArrayRef =
1558+
Arc::new(StringArray::from(vec![Some("x"), None, Some("zzz")]));
1559+
let replacement =
1560+
dictionary_array(vec![Some(0), Some(2), None, Some(1)], replacement_values);
1561+
let mut expected = vec![0; replacement.len()];
1562+
create_hashes([&replacement], &random_state, &mut expected)?;
1563+
let mut actual = vec![0; replacement.len()];
1564+
cache.create_hashes(std::slice::from_ref(&replacement), &mut actual)?;
1565+
assert_eq!(actual, expected);
1566+
assert_eq!(cache.entries[0].rebuilds, 1);
1567+
assert_eq!(cache.size(), 0);
1568+
cache.create_hashes(&[replacement], &mut actual)?;
1569+
assert_eq!(actual, expected);
1570+
assert_eq!(cache.entries[0].rebuilds, 2);
1571+
Ok(())
1572+
}
1573+
1574+
#[cfg(not(feature = "force_hash_collisions"))]
1575+
#[test]
1576+
fn dictionary_hash_cache_matches_composite_hashes() -> Result<()> {
1577+
let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"]));
1578+
let dictionary = dictionary_array(vec![Some(2), Some(0), Some(1), None], values);
1579+
let primitive: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
1580+
let arrays = [primitive, dictionary];
1581+
let random_state = RandomState::with_seed(42);
1582+
let mut expected = vec![0; arrays[0].len()];
1583+
create_hashes(&arrays, &random_state, &mut expected)?;
1584+
1585+
let mut actual = vec![0; arrays[0].len()];
1586+
let mut cache = DictionaryHashCache::new(random_state);
1587+
cache.create_hashes(&arrays, &mut actual)?;
1588+
assert_eq!(actual, expected);
1589+
assert_eq!(cache.entries[1].rebuilds, 0);
1590+
assert_eq!(cache.size(), 0);
1591+
cache.create_hashes(&arrays, &mut actual)?;
1592+
assert_eq!(actual, expected);
1593+
assert_eq!(cache.entries[1].rebuilds, 1);
1594+
assert!(cache.size() > 0);
1595+
Ok(())
1596+
}
1597+
13561598
#[cfg(not(feature = "force_hash_collisions"))]
13571599
#[derive(Default)]
13581600
struct TestHasher(u64);

0 commit comments

Comments
 (0)