jayzhan211 commented on code in PR #26816:
URL: https://github.com/apache/datafusion/pull/26816#discussion_r4499105695
##########
datafusion/common/src/hash_utils.rs:
##########
@@ -1357,6 +1489,187 @@ 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 {
+ CachingArrayHasher::new(columns.len(), 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
Review Comment:
These pointer checks also pass on a miss because the `Vec` is recycled, so
nothing here shows the rehash was skipped. I made the cache never hit and made
`BatchPartitioner` ignore the hasher, and all the tests still passed in both
cases. Since reuse is the point of this PR and #26814 asks for it to be tested
directly, could you add this in this PR? It poisons the cached hashes, so only
a real hit can return the sentinel, and it fails when the cache never hits. It
also covers equal values behind a new `Arc` and a dropped `Arc`:
```rs
#[cfg(not(feature = "force_hash_collisions"))]
#[test]
fn caching_array_hasher_skips_rehash_only_for_same_values_arc() ->
Result<()> {
let random_state = RandomState::with_seed(0);
let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
let batch = vec![string_dictionary(vec![0, 1], &values)];
let mut hasher = caching_hasher(&batch, &random_state);
let mut hashes = vec![0; 2];
hasher.create_hashes(&batch, &mut hashes)?;
let fresh = hashes.clone();
// Poison the cached hashes: only a cache hit can return the sentinel.
const SENTINEL: u64 = 0xDEAD_BEEF;
hasher.dictionaries[0].as_mut().unwrap().1.fill(SENTINEL);
let same_arc = vec![string_dictionary(vec![1, 0], &values)];
hasher.create_hashes(&same_arc, &mut hashes)?;
assert_eq!(hashes, vec![SENTINEL; 2]);
// Equal contents behind a different `Arc` is a miss.
let equal_values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
let new_arc = vec![string_dictionary(vec![0, 1], &equal_values)];
hasher.create_hashes(&new_arc, &mut hashes)?;
assert_eq!(hashes, fresh);
// A dropped values array must not be matched by a later allocation.
hasher.dictionaries[0].as_mut().unwrap().1.fill(SENTINEL);
drop(new_arc);
drop(equal_values);
for _ in 0..1000 {
let reused: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
let batch = vec![string_dictionary(vec![0, 1], &reused)];
hasher.create_hashes(&batch, &mut hashes)?;
assert_eq!(hashes, fresh);
hasher.dictionaries[0].as_mut().unwrap().1.fill(SENTINEL);
}
Ok(())
}
```
The pointer assertions can go once this is in. The test in my comment on
`size()` covers the `BatchPartitioner` wiring.
##########
datafusion/common/src/hash_utils.rs:
##########
@@ -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.
+ ///
+ /// 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(num_columns: usize, random_state: RandomState) -> Self {
+ Self {
+ random_state,
+ dictionaries: vec![None; num_columns],
+ }
+ }
+
+ /// 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 {
Review Comment:
Nothing calls `size()`, so this cache isn't in the memory pool, and #26814
asks for it to be visible in accounting. It used to be per-batch scratch for
one column at a time; now it's 8 bytes per value across all dictionary key
columns, held per input partition for the whole stream. Could you reserve it in
this PR? This version gives up the cache when the pool refuses, so a query
never fails where it wouldn't have before. I ran it: the repartition and
`memory_limit` tests pass, and the new test fails if the partitioner stops
using the hasher.
```diff
--- a/datafusion/common/src/hash_utils.rs
+++ b/datafusion/common/src/hash_utils.rs
+ /// Drops all cached value hashes. Later batches rehash and refill the
cache.
+ pub fn clear(&mut self) {
+ self.dictionaries.fill(None);
+ }
+
/// Returns the allocated size of the cache.
pub fn size(&self) -> usize {
--- a/datafusion/physical-plan/src/repartition/mod.rs
+++ b/datafusion/physical-plan/src/repartition/mod.rs
-use datafusion_execution::memory_pool::MemoryConsumer;
+use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
@@ spawning pull_from_input
num_input_partitions,
+ MemoryConsumer::new(format!("{name} hash cache[{i}]"))
+ .register(context.memory_pool()),
));
@@ impl BatchPartitioner
+ /// Accounts the retained dictionary hash cache in `reservation`. If the
+ /// pool refuses it, drops the cache so later batches pay the uncached
cost.
+ fn reserve_hash_cache(&mut self, reservation: &MemoryReservation) {
+ if let BatchPartitionerState::Hash { hasher, .. } = &mut self.state
+ && reservation.try_resize(hasher.size()).is_err()
+ {
+ hasher.clear();
+ reservation.resize(hasher.size());
+ }
+ }
@@ pull_from_input
+ #[expect(clippy::too_many_arguments)]
async fn pull_from_input(
@@
num_input_partitions: usize,
+ hash_cache_reservation: MemoryReservation,
) -> Result<()> {
@@ after the partition_iter loop
timer.done();
}
+ partitioner.reserve_hash_cache(&hash_cache_reservation);
```
```rs
#[test]
fn hash_partitioner_caches_and_reserves_dictionary_hashes() -> Result<()> {
use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool};
let dictionary_type =
DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8));
let schema = Arc::new(Schema::new(vec![Field::new("key",
dictionary_type, true)]));
let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb",
"ccc"]));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(DictionaryArray::<Int32Type>::try_new(
Int32Array::from(vec![0, 1, 2]),
Arc::clone(&values),
)?)],
)?;
let cached = |partitioner: &BatchPartitioner| {
let BatchPartitionerState::Hash { hasher, .. } = &partitioner.state
else {
unreachable!()
};
(hasher.size(), hasher.dictionary_hashes(0).map(|(v, _)| v))
};
// The pool has room: the cache is kept and reserved.
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024));
let reservation = MemoryConsumer::new("cache").register(&pool);
let mut partitioner = BatchPartitioner::new_hash_partitioner(
vec![col("key", &schema)?],
4,
metrics::Time::new(),
)?;
partitioner
.partition_iter(batch.clone())?
.collect::<Result<Vec<_>>>()?;
partitioner.reserve_hash_cache(&reservation);
let (size, cached_values) = cached(&partitioner);
assert!(Arc::ptr_eq(&cached_values.unwrap(), &values));
assert!(size >= 3 * size_of::<u64>());
assert_eq!(pool.reserved(), size);
// The pool has no room: the cache is dropped, partitioning still works.
let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(8));
let reservation = MemoryConsumer::new("cache").register(&pool);
partitioner.reserve_hash_cache(&reservation);
assert!(cached(&partitioner).1.is_none());
partitioner
.partition_iter(batch)?
.collect::<Result<Vec<_>>>()?;
Ok(())
}
```
##########
datafusion/common/src/hash_utils.rs:
##########
@@ -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.
+ ///
+ /// 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(num_columns: usize, random_state: RandomState) -> Self {
+ Self {
+ random_state,
+ dictionaries: vec![None; num_columns],
+ }
+ }
+
+ /// 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);
Review Comment:
Optional cleanup: clearing first means each slot is only written once, and
the randomized comparison against `create_hashes` still passes.
```suggestion
// Fresh hashing leaves null value slots unchanged.
hashes.clear();
hashes.resize(values.len(), 0);
```
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -20,6 +20,8 @@ use super::{
};
use crate::coalesce::LimitedBatchCoalescer;
use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
+use crate::hash_utils::CachingArrayHasher;
+#[cfg(test)]
use crate::hash_utils::create_hashes;
Review Comment:
Optional cleanup: only the test at the bottom of this file uses this import,
so it can move into `mod tests` and drop the `cfg`. Clippy passes with this.
```diff
-#[cfg(test)]
-use crate::hash_utils::create_hashes;
@@
mod tests {
+ use crate::hash_utils::create_hashes;
```
##########
datafusion/common/src/hash_utils.rs:
##########
@@ -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.
Review Comment:
Optional cleanup: the address comparison is only safe because the `Weak`
also keeps the allocation reserved, which is worth stating so nobody later
swaps it for a raw pointer.
```suggestion
/// One entry per column. Weak identity avoids retaining dictionary
values,
/// and keeps their allocation reserved so a new `Arc` can't reuse the
address.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
