jayzhan211 commented on code in PR #26816:
URL: https://github.com/apache/datafusion/pull/26816#discussion_r4502204898
##########
datafusion/common/src/hash_utils.rs:
##########
@@ -28,6 +32,140 @@ 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,
+ /// and keeps their allocation reserved so a new `Arc` can't reuse the
address.
+ 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, _)| {
Review Comment:
A cache miss hashes every value, which skips the per-row bound #26703 added
to `hash_dictionary`. So with more than 64 values per row and a changing values
`Arc`, this is slower than main: 1.60 s to 6.67 s (4.17x) on a `LIMIT` subquery
that interleaves 3 files with 2M values each. Could you leave those misses to
`hash_dictionary`, after rebasing? With this change that query takes 1.64 s,
and the 0.83x win on a shared `Arc` below the bound stays. Shared values above
the bound stop being cached, which also caps the retained buffer at 4 MiB per
column for 8192-row batches.
```diff
-use crate::utils::compact_oversized_dictionary;
+use crate::utils::{MAX_DICTIONARY_VALUES_PER_ROW,
compact_oversized_dictionary};
@@ CachingArrayHasher::hash_array
- if !dictionary.as_ref().is_some_and(|(cached, _)| {
+ let cached = dictionary.as_ref().is_some_and(|(cached,
_)| {
std::ptr::addr_eq(cached.as_ptr(),
Arc::as_ptr(values))
- }) {
+ });
+ // Above the bound `hash_dictionary` only hashes the
values
+ // these rows reference, so a miss must not hash all of
them.
+ if !cached && values.len() >
MAX_DICTIONARY_VALUES_PER_ROW * array.len()
+ {
+ hash_single_array(
+ array,
+ &self.random_state,
+ hashes_buffer,
+ rehash,
+ )?;
+ return Ok(());
+ }
+ if !cached {
```
<details>
<summary>Boundary test and the SQL repro</summary>
```rs
#[cfg(not(feature = "force_hash_collisions"))]
#[test]
fn caching_array_hasher_does_not_fill_above_values_per_row_bound() ->
Result<()> {
let random_state = RandomState::with_seed(0);
let dictionary_of = |len: usize| -> ArrayRef {
let values: ArrayRef = Arc::new(StringArray::from_iter_values(
(0..len).map(|i| format!("{i}")),
));
string_dictionary(vec![0], &values)
};
// One row may hash up to the bound of values: this fills the cache.
let at_bound = vec![dictionary_of(MAX_DICTIONARY_VALUES_PER_ROW)];
let mut hasher = caching_hasher(&at_bound, &random_state);
assert_hashes_match(&mut hasher, &at_bound, &random_state)?;
assert!(hasher.dictionary_hashes(0).is_some());
// One value more is a miss that must not hash, or retain, all values.
let above_bound = vec![dictionary_of(MAX_DICTIONARY_VALUES_PER_ROW + 1)];
assert_hashes_match(&mut hasher, &above_bound, &random_state)?;
let cached_len = hasher.dictionaries[0].as_ref().unwrap().1.len();
assert_eq!(cached_len, MAX_DICTIONARY_VALUES_PER_ROW);
Ok(())
}
```
Write this COPY three times, to `b/f0.parquet`, `b/f1.parquet` and
`b/f2.parquet`:
```sql
COPY (SELECT 'value_' || CAST(floor(random()*2000000) AS BIGINT) AS k, value
AS id
FROM generate_series(1,8000000))
TO 'b/f0.parquet' STORED AS PARQUET OPTIONS(
'format.dictionary_page_size_limit' '200000000',
'format.max_row_group_size' '100000000',
'format.dictionary_enabled' 'true',
'format.compression' 'uncompressed');
SET datafusion.execution.parquet.enable_rle_to_dictionary = true;
SET datafusion.execution.target_partitions = 12;
SET datafusion.optimizer.hash_join_single_partition_threshold = 0;
SET datafusion.optimizer.hash_join_single_partition_threshold_rows = 0;
SET datafusion.optimizer.enable_join_dynamic_filter_pushdown = false;
CREATE EXTERNAL TABLE tb STORED AS PARQUET LOCATION 'b/';
SELECT COUNT(*) FROM (SELECT k FROM tb LIMIT 24000000) a
JOIN (SELECT DISTINCT k FROM tb WHERE id = 5) b ON a.k = b.k;
```
Min of 4 interleaved runs, same result (43) on all sides. Hash time alone
per 8192-row batch, 8-byte values, main vs a miss on every batch: 65 values per
row 78 vs 399 us, 512 per row 95 vs 3068 us.
</details>
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2832,6 +2852,57 @@ mod tests {
);
}
+ #[test]
Review Comment:
This test panics under `force_hash_collisions`, because the cache is never
filled there, so the merge queue's hash-collisions job would fail. That one's
on me, since it came from my suggestion.
```suggestion
#[cfg(not(feature = "force_hash_collisions"))]
#[test]
```
##########
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:
Correction to my earlier comment: the reservation itself can't fail a query,
but it does take pool room, so another operator can hit the limit slightly
sooner. In a tight 2.1 GB run, the same join passed 2 of 6 times on main and 0
of 6 with the cache, and at 2.5 GB and above each side passed 5 or 6 of 6.
Giving up the cache on refusal is still right, so there's nothing to change.
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2343,6 +2361,7 @@ impl RepartitionExec {
}
timer.done();
}
+ partitioner.reserve_hash_cache(&hash_cache_reservation);
Review Comment:
Nothing tests this call: I removed it and all 80 repartition tests still
pass, so the cache could silently fall out of the memory pool again. Could you
add this test in this PR? It fails with the call removed and passes otherwise.
<details>
<summary>End-to-end test that records the hash cache reservation</summary>
```rs
#[cfg(not(feature = "force_hash_collisions"))]
#[tokio::test]
async fn hash_repartition_reserves_dictionary_hash_cache() -> Result<()> {
use datafusion_execution::memory_pool::{
MemoryPool, MemoryReservation, UnboundedMemoryPool,
};
/// Records the largest reservation of the hash cache consumers.
#[derive(Debug, Default)]
struct CachePeakPool {
inner: UnboundedMemoryPool,
cache_peak: AtomicUsize,
}
impl CachePeakPool {
fn record(&self, reservation: &MemoryReservation, additional: usize)
{
if reservation.consumer().name().contains("hash cache") {
self.cache_peak.fetch_max(
reservation.size() + additional,
AtomicOrdering::Relaxed,
);
}
}
}
impl Display for CachePeakPool {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str("CachePeakPool")
}
}
impl MemoryPool for CachePeakPool {
fn name(&self) -> &str {
"cache_peak"
}
fn grow(&self, reservation: &MemoryReservation, additional: usize) {
self.record(reservation, additional);
self.inner.grow(reservation, additional);
}
fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
self.inner.shrink(reservation, shrink);
}
fn try_grow(
&self,
reservation: &MemoryReservation,
additional: usize,
) -> Result<()> {
self.record(reservation, additional);
self.inner.try_grow(reservation, additional)
}
fn reserved(&self) -> usize {
self.inner.reserved()
}
}
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]),
values,
)?)],
)?;
let pool = Arc::new(CachePeakPool::default());
let runtime = RuntimeEnvBuilder::new()
.with_memory_pool(Arc::clone(&pool) as Arc<dyn MemoryPool>)
.build_arc()?;
let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime));
let exec = TestMemoryExec::try_new_exec(&[vec![batch]],
Arc::clone(&schema), None)?;
let exec = RepartitionExec::try_new(
exec,
Partitioning::Hash(vec![col("key", &schema)?], 1),
)?;
let output = crate::common::collect(exec.execute(0, task_ctx)?).await?;
assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
assert!(pool.cache_peak.load(AtomicOrdering::Relaxed) >= 3 *
size_of::<u64>());
Ok(())
}
```
</details>
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2832,6 +2852,57 @@ mod tests {
);
}
+ #[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());
Review Comment:
Optional cleanup: this also checks that the reservation shrinks after the
cache is dropped, which no test covers today. It fails if the
`reservation.resize` call goes away.
```suggestion
assert!(cached(&partitioner).1.is_none());
assert_eq!(pool.reserved(), cached(&partitioner).0);
```
--
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]
