jayzhan211 commented on code in PR #26826:
URL: https://github.com/apache/datafusion/pull/26826#discussion_r4505351464
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs:
##########
@@ -298,9 +334,76 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
group_values,
hashes_buffer: Default::default(),
random_state: crate::aggregates::AGGREGATION_HASH_SEED,
+ has_dictionary: !STREAMING && has_dictionary,
+ column_hashes_buffer: Default::default(),
})
}
+ #[cfg(not(feature = "force_hash_collisions"))]
+ fn create_dictionary_hashes(
+ &mut self,
+ cols: &[ArrayRef],
+ hashes: &mut [u64],
+ ) -> Result<()> {
+ let group_unique = self.group_values.len() == 1;
+
+ let use_cache = cols.iter().any(|array| {
+ array.as_any_dictionary_opt().is_some_and(|dictionary| {
+ let values = dictionary.values();
+ values.len() <=
MAX_DICTIONARY_VALUES_PER_ROW.saturating_mul(array.len())
+ && (group_unique || values.len() > array.len())
+ })
+ });
+ if !use_cache {
Review Comment:
@Rich-T-kid thanks for this! Batches that take the early `create_hashes`
return and batches that take the cached loop hash the same group differently,
so `GROUP BY dict_col, int_col` returns duplicate groups. The loop folds a
later column in with `combine_hashes`, but `create_hashes` feeds the previous
hash into a seeded hasher for primitive and view columns and leaves null rows
unchanged. The loop only runs when a dictionary has more values than rows, so a
filter or a short last batch mixes the two paths. This query returns `a 1 1`
twice on this branch and `a 1 2` on main:
```sql
set datafusion.execution.target_partitions = 1;
set datafusion.execution.batch_size = 4;
CREATE TABLE t AS
SELECT arrow_cast(chr(97 + value % 4), 'Dictionary(Int32, Utf8)') AS d, 1 AS
i, value AS v
FROM range(8);
SELECT d, i, count(*) FROM t WHERE v < 5 GROUP BY d, i ORDER BY d, i;
```
Hashing every batch through one loop, with each column hashed exactly as
`create_hashes` would, fixes it and also removes `column_hashes_buffer`. Could
you apply this and add the test below?
```diff
--- a/datafusion/common/src/hash_utils.rs
+++ b/datafusion/common/src/hash_utils.rs
@@ pub fn scatter_dictionary_hashes<K: ArrowDictionaryKeyType>(
Ok(())
}
+/// Hashes one column into `hashes_buffer` exactly as [`create_hashes`] does
+/// for that column; `rehash` is `true` for every column after the first.
+///
+/// # Public Only for Internal Use:
+///
+/// This is not a public API and is for internal use only; see [API policy]
for details.
+///
+/// [API policy]:
https://datafusion.apache.org/contributor-guide/api-health.html#internal-public-apis
+#[doc(hidden)]
+#[cfg(not(feature = "force_hash_collisions"))]
+pub fn hash_column(
+ array: &dyn Array,
+ random_state: &impl HashState,
+ hashes_buffer: &mut [u64],
+ rehash: bool,
+) -> Result<()> {
+ hash_single_array(array, random_state, hashes_buffer, rehash)
+}
+
---
a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
+++
b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
-use datafusion_common::hash_utils::{combine_hashes,
scatter_dictionary_hashes};
+use datafusion_common::hash_utils::{hash_column, scatter_dictionary_hashes};
@@ (struct field, `try_new`, `size`, `clear_shrink`, emit shrink)
- column_hashes_buffer: Vec<u64>,
- column_hashes_buffer: Default::default(),
- + self.column_hashes_buffer.allocated_size()
- self.column_hashes_buffer.clear();
- self.column_hashes_buffer.shrink_to_fit();
- self.column_hashes_buffer.clear();
- self.column_hashes_buffer.shrink_to(num_rows);
@@ fn create_dictionary_hashes(
let group_unique = self.group_values.len() == 1;
-
- let use_cache = cols.iter().any(|array| {
- array.as_any_dictionary_opt().is_some_and(|dictionary| {
- let values = dictionary.values();
- values.len() <=
MAX_DICTIONARY_VALUES_PER_ROW.saturating_mul(array.len())
- && (group_unique || values.len() > array.len())
- })
- });
- if !use_cache {
- create_hashes(cols, &crate::aggregates::AGGREGATION_HASH_SEED,
hashes)?;
- return Ok(());
- }
-
for (index, (array, group_value)) in
cols.iter().zip(&mut self.group_values).enumerate()
{
+ let rehash = index != 0;
let cached =
array.as_any_dictionary_opt().and_then(|dictionary| {
@@
- index != 0,
+ rehash,
)?,
_ => unreachable!(),
}
- } else if index == 0 {
- create_hashes(std::slice::from_ref(array),
&self.random_state, hashes)?;
} else {
- self.column_hashes_buffer.clear();
- self.column_hashes_buffer.resize(hashes.len(), 0);
- create_hashes(
- std::slice::from_ref(array),
- &self.random_state,
- &mut self.column_hashes_buffer,
- )?;
- for (hash, column_hash) in
- hashes.iter_mut().zip(&self.column_hashes_buffer)
- {
- *hash = combine_hashes(*column_hash, *hash);
- }
+ hash_column(array.as_ref(), &self.random_state, hashes,
rehash)?;
}
```
<details><summary>Unit test: same group through both hash paths (passes on
main, fails on this branch, passes with the fix)</summary>
```rs
#[test]
fn dictionary_hash_paths_agree_on_existing_groups() {
let schema = Arc::new(Schema::new(vec![
Field::new(
"d",
DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Utf8)),
true,
),
Field::new("i", DataType::Int32, true),
]));
let mut gv = GroupValuesColumn::<false>::try_new(schema).unwrap();
let mut groups = vec![];
// 2 values <= 4 rows: hashed without the dictionary cache
let small: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
let d1: ArrayRef = Arc::new(DictionaryArray::<Int8Type>::new(
Int8Array::from(vec![0, 1, 0, 1]),
small,
));
let i1: ArrayRef =
Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(3)]));
gv.intern(&[d1, i1], &mut groups).unwrap();
assert_eq!(groups, [0, 1, 2, 3]);
// 3 values > 2 rows: hashed through the dictionary cache
let big: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
let d2: ArrayRef = Arc::new(DictionaryArray::<Int8Type>::new(
Int8Array::from(vec![0, 0]),
big,
));
let i2: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None]));
gv.intern(&[d2, i2], &mut groups).unwrap();
// ('a', 1) is group 0 and ('a', NULL) is group 2
assert_eq!(groups, [0, 2]);
assert_eq!(gv.len(), 4);
}
```
</details>
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs:
##########
@@ -598,7 +649,12 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
self.value_dedup_size = 0;
self.null_inner_slot = None;
- self.hash_values(&all_inner_values);
+ if self.values_are_group_unique {
Review Comment:
@Rich-T-kid the direct key-to-group mapping returns wrong groups after
`EmitTo::First` (partial aggregation under memory pressure emits early this
way) when a NULL group remains. `take_n` moves the null slot to the end, so
inner slots stop matching group indices. With groups `'a' -> 0, NULL -> 1, 'b'
-> 2`, `emit(EmitTo::First(1))` and then two batches with keys `['b', NULL]` on
the same values array give `[1, 0]` from the normal path but `[0, 1]` from the
direct mapping, so the NULL and `'b'` aggregates get swapped.
Making the reorder depend on `values_are_group_unique` isn't enough: when
the first batches carry oversized dictionaries, the flag is still false when
`take_n` runs. Nothing needs the null slot last any more (`into_dict` remaps
keys and `check_key_overflow` only counts slots), so could you drop the reorder
(lines 625-647) and add the two tests below?
```diff
- // Null deferred to last so null_inner_slot is always the highest
index
- // and check_key_overflow can subtract it without a false overflow.
+ // Remaining groups keep their order, so for a sole group column
each
+ // inner slot stays equal to its group index, which
+ // `try_map_dictionary_groups` relies on.
let remaining = self.group_to_inner[n..].to_vec();
let mut old_to_new = vec![usize::MAX; old_inner_len];
let mut new_to_old: Vec<usize> = Vec::new();
- let mut null_old_slot: Option<usize> = None;
for &old in &remaining {
- if all_inner_values.is_null(old) {
- if null_old_slot.is_none() {
- null_old_slot = Some(old);
- }
- continue;
- }
if old_to_new[old] == usize::MAX {
old_to_new[old] = new_to_old.len();
new_to_old.push(old);
}
}
- if let Some(old) = null_old_slot {
- old_to_new[old] = new_to_old.len();
- new_to_old.push(old);
- }
```
<details><summary>Two unit tests: emit with a NULL group, and the same after
oversized dictionaries (both pass on main, fail on this branch, pass with the
fix)</summary>
```rs
#[test]
fn single_dictionary_direct_mapping_after_emit_first_with_null_group() {
let schema = Arc::new(Schema::new(vec![Field::new(
"d",
DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Utf8)),
true,
)]));
let mut gv = GroupValuesColumn::<false>::try_new(schema).unwrap();
let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"]));
let dict = |keys: Vec<Option<i8>>| -> ArrayRef {
Arc::new(DictionaryArray::<Int8Type>::new(
Int8Array::from(keys),
Arc::clone(&values),
))
};
let mut groups = vec![];
// 'a' -> 0, NULL -> 1, 'b' -> 2
gv.intern(&[dict(vec![Some(0), None, Some(1)])], &mut groups)
.unwrap();
assert_eq!(groups, [0, 1, 2]);
// emit 'a'; remaining groups: NULL -> 0, 'b' -> 1
gv.emit(EmitTo::First(1)).unwrap();
gv.intern(&[dict(vec![Some(1), None])], &mut groups).unwrap();
assert_eq!(groups, [1, 0]);
// same values array again: direct mapping
gv.intern(&[dict(vec![Some(1), None])], &mut groups).unwrap();
assert_eq!(groups, [1, 0]);
}
#[test]
fn single_dictionary_direct_mapping_after_oversized_batches_and_emit_first()
{
let schema = Arc::new(Schema::new(vec![Field::new(
"d",
DataType::Dictionary(Box::new(DataType::Int16),
Box::new(DataType::Utf8)),
true,
)]));
let mut gv = GroupValuesColumn::<false>::try_new(schema).unwrap();
let mut strs: Vec<String> = (0..200).map(|i| format!("v{i}")).collect();
strs[0] = "a".into();
strs[1] = "b".into();
let big: ArrayRef = Arc::new(StringArray::from(strs));
let dict = |keys: Vec<Option<i16>>, values: &ArrayRef| -> ArrayRef {
Arc::new(DictionaryArray::<arrow::datatypes::Int16Type>::new(
arrow::array::Int16Array::from(keys),
Arc::clone(values),
))
};
let mut groups = vec![];
// 200 values > 64 * 3 rows: oversized, hashed without the cache
gv.intern(&[dict(vec![Some(0), None, Some(1)], &big)], &mut groups)
.unwrap();
assert_eq!(groups, [0, 1, 2]);
gv.emit(EmitTo::First(1)).unwrap(); // NULL -> 0, 'b' -> 1
let small: ArrayRef = Arc::new(StringArray::from(vec!["b"]));
gv.intern(&[dict(vec![Some(0), None], &small)], &mut groups)
.unwrap();
assert_eq!(groups, [1, 0]);
gv.intern(&[dict(vec![Some(0), None], &small)], &mut groups)
.unwrap();
assert_eq!(groups, [1, 0]);
}
```
</details>
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs:
##########
@@ -331,13 +347,56 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
Some(val_idx) => {
let dict_values = dict.values();
self.sync_value_cache(dict_values);
- self.find_or_insert_value(dict_values, val_idx,
self.val_hashes[val_idx])?
+ self.insert_value(dict_values, val_idx)?
}
};
self.group_to_inner.push(inner_slot);
self.check_key_overflow()
}
+ fn dictionary_hashes<'a>(
+ &'a mut self,
+ values: &ArrayRef,
+ group_unique: bool,
+ ) -> Option<&'a [u64]> {
+ self.values_are_group_unique = group_unique;
+ self.sync_value_cache(values);
+ Some(&self.val_hashes)
+ }
+
+ fn try_map_dictionary_groups(
+ &self,
+ array: &ArrayRef,
+ groups: &mut Vec<usize>,
+ ) -> bool {
+ if !self.values_are_group_unique {
+ return false;
+ }
+
+ let dictionary = array.as_dictionary::<K>();
+ if !self.is_cached(dictionary.values()) {
+ return false;
+ }
+
+ let keys = dictionary.keys();
+ let raw_keys = keys.values();
+ let nulls = keys.nulls();
+ groups.clear();
+ for row in 0..dictionary.len() {
+ let group = if nulls.is_some_and(|nulls| nulls.is_null(row)) {
+ self.null_inner_slot.unwrap_or(usize::MAX)
+ } else {
+ self.val_to_inner[raw_keys[row].as_usize()]
+ };
+ if group == usize::MAX {
Review Comment:
@Rich-T-kid no test reaches this fallback: removing the `usize::MAX` check
still passes every test, yet without it a batch that shares the values array
and brings a new value returns group `usize::MAX`. That is the common case for
a shared Parquet dictionary, so could you add this test?
<details><summary>Unit test: shared values array, second batch adds a new
value (passes now, fails if the fallback is removed)</summary>
```rs
#[test]
#[cfg(not(feature = "force_hash_collisions"))]
fn single_dictionary_shared_values_new_key_falls_back() {
let schema = Arc::new(Schema::new(vec![Field::new(
"category",
DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Utf8)),
true,
)]));
let mut group_values =
GroupValuesColumn::<false>::try_new(schema).unwrap();
let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
let dictionary = |keys: Vec<i8>| -> ArrayRef {
Arc::new(DictionaryArray::<Int8Type>::new(
Int8Array::from(keys),
Arc::clone(&values),
))
};
let mut groups = vec![];
group_values.intern(&[dictionary(vec![0, 1])], &mut groups).unwrap();
assert_eq!(groups, [0, 1]);
// Same values array, but "c" is not a group yet: the direct mapping
// must fall back and create group 2.
group_values
.intern(&[dictionary(vec![1, 2, 0])], &mut groups)
.unwrap();
assert_eq!(groups, [1, 2, 0]);
assert_eq!(group_values.len(), 3);
}
```
</details>
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs:
##########
@@ -2124,6 +2250,137 @@ mod tests {
assert_eq!(out[1].as_dictionary::<Int8Type>().values().len(), 2);
}
+ #[test]
+ #[cfg(not(feature = "force_hash_collisions"))]
+ fn multi_dictionary_hash_cache_tracks_values() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("label", DataType::Int32, false),
+ Field::new(
+ "category",
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Utf8)),
+ true,
+ ),
+ ]));
+ let dictionary = |keys, values| -> ArrayRef {
+ Arc::new(DictionaryArray::<Int8Type>::new(
+ Int8Array::from(keys),
+ values,
+ ))
+ };
+ let first_values: ArrayRef = Arc::new(StringArray::from(vec![
+ Some("cat"),
+ Some("dog"),
+ None,
+ Some("owl"),
+ ]));
+ let first = vec![
+ Arc::new(Int32Array::from(vec![0, 1])) as ArrayRef,
+ dictionary(vec![Some(0), Some(1)], Arc::clone(&first_values)),
+ ];
+
+ let mut group_values =
GroupValuesColumn::<false>::try_new(schema).unwrap();
+ let mut groups = vec![];
+ group_values.intern(&first, &mut groups).unwrap();
+ assert_eq!(groups, [0, 1]);
+ let allocation = group_values.group_values[1]
+ .dictionary_hashes(&first_values, false)
+ .unwrap()
+ .as_ptr();
+
+ let reordered = vec![
+ Arc::new(Int32Array::from(vec![0, 1])) as ArrayRef,
+ dictionary(vec![Some(1), Some(0)], Arc::clone(&first_values)),
+ ];
+ group_values.intern(&reordered, &mut groups).unwrap();
+ assert_eq!(groups, [2, 3]);
+ assert_eq!(
+ group_values.group_values[1]
+ .dictionary_hashes(&first_values, false)
+ .unwrap()
+ .as_ptr(),
+ allocation
Review Comment:
Optional cleanup: this pointer check can't fail. `hash_values` clears and
refills `val_hashes` in place, so the pointer stays the same even when the
cache is rebuilt, which is exactly what the "changed values" step below
asserts. To show reuse, assert something that changes on a rebuild, or count
calls to `hash_values` in a test-only counter.
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs:
##########
@@ -298,9 +334,76 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
group_values,
hashes_buffer: Default::default(),
random_state: crate::aggregates::AGGREGATION_HASH_SEED,
+ has_dictionary: !STREAMING && has_dictionary,
+ column_hashes_buffer: Default::default(),
})
}
+ #[cfg(not(feature = "force_hash_collisions"))]
+ fn create_dictionary_hashes(
Review Comment:
@Rich-T-kid the posted benchmarks only cover one dictionary column with one
shared values array, the best case for both new paths. Could you add these
cases to `dictionary_group_values.rs` and post the bot results, after the hash
fix above since it changes the multi-column path?
- two values arrays alternating between batches (one dictionary column, and
dictionary + `Int32`), cardinality 100k
- dictionary + `Int32` with one shared values array larger than the batch
(32k and 100k values)
- a fresh 20-value dictionary per batch (e.g. a cast in a projection)
My laptop runs were too noisy to use (other jobs ran during them), but the
alternating multi-column case and the fresh-dictionary case came out slower on
both runs.
--
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]
