jayzhan211 commented on code in PR #26966:
URL: https://github.com/apache/datafusion/pull/26966#discussion_r4502193988
##########
datafusion/common/src/hash_utils.rs:
##########
@@ -563,19 +563,23 @@ fn hash_dictionary_scatter<
hashes_buffer: &mut [u64],
) {
let dict_values = array.values();
- if HAS_NULL_KEYS {
- for (hash, key) in hashes_buffer.iter_mut().zip(array.keys().iter()) {
- if let Some(key) = key {
- let idx = key.as_usize();
- if !HAS_NULL_VALUES || dict_values.is_valid(idx) {
- // SAFETY: `idx` is a valid Arrow dict key, in-bounds for
`dict_hashes`.
- let dict_hash = unsafe { *dict_hashes.get_unchecked(idx) };
- *hash = if MULTI_COL {
- combine_hashes(dict_hash, *hash)
- } else {
- dict_hash
- };
- }
+ if HAS_NULL_KEYS && let Some(key_nulls) = array.keys().nulls() {
+ // Visit only valid key slots: null slots keep their hash, and their
+ // (possibly out-of-range) key values are never read
+ let keys = array.keys().values();
+ // Like the `zip` in the loop below, stop at the end of `hashes_buffer`
+ let len = hashes_buffer.len();
+ for i in key_nulls.valid_indices().take_while(|&i| i < len) {
Review Comment:
This still visits mostly-valid rows one set bit at a time, so with 1 null in
8,192 rows `create_hashes` takes 5.5 µs, against 2.6 µs with the 64-row word
loop from #26848 (3 columns: 18.6 µs, against 18.0 µs on main and 10.3 µs with
the word loop). Your ±4–9% comparison used 3%, 50% and 97% nulls, where few or
no words are fully valid, so it couldn't show this gap. Could you switch to the
loop below in this PR? Your tests pass on it unchanged, the no-null cases don't
move, and its worst case against this PR's loop is 5% slower (97% nulls, 3
columns).
```diff
@@ -563,37 +563,52 @@ fn hash_dictionary_scatter<
hashes_buffer: &mut [u64],
) {
let dict_values = array.values();
+ let scatter_row = |hash: &mut u64, key: &K::Native| {
+ let idx = key.as_usize();
+ if !HAS_NULL_VALUES || dict_values.is_valid(idx) {
+ // SAFETY: `idx` is a valid Arrow dict key, in-bounds for
`dict_hashes`.
+ let dict_hash = unsafe { *dict_hashes.get_unchecked(idx) };
+ *hash = if MULTI_COL {
+ combine_hashes(dict_hash, *hash)
+ } else {
+ dict_hash
+ };
+ }
+ };
+ let keys = array.keys().values();
if HAS_NULL_KEYS && let Some(key_nulls) = array.keys().nulls() {
- // Visit only valid key slots: null slots keep their hash, and their
- // (possibly out-of-range) key values are never read
- let keys = array.keys().values();
- // Like the `zip` in the loop below, stop at the end of
`hashes_buffer`
- let len = hashes_buffer.len();
- for i in key_nulls.valid_indices().take_while(|&i| i < len) {
- let idx = keys[i].as_usize();
- if !HAS_NULL_VALUES || dict_values.is_valid(idx) {
- // SAFETY: `idx` is a valid Arrow dict key, in-bounds for
`dict_hashes`.
- let dict_hash = unsafe { *dict_hashes.get_unchecked(idx) };
- let hash = &mut hashes_buffer[i];
- *hash = if MULTI_COL {
- combine_hashes(dict_hash, *hash)
- } else {
- dict_hash
- };
+ // Like the `zip` below, stop at the end of `hashes_buffer`. Null
slots keep
+ // their hash, and their (possibly out-of-range) keys are never
read.
+ let len = hashes_buffer.len().min(keys.len());
+ let (hashes, keys) = (&mut hashes_buffer[..len], &keys[..len]);
+ let chunks = key_nulls.inner().bit_chunks();
+ let words = chunks
+ .iter()
+ .chain(std::iter::once(chunks.remainder_bits()));
+ for (word_idx, mut word) in words.enumerate() {
+ let start = word_idx * 64;
+ if start >= len {
+ break;
+ }
+ if word == u64::MAX {
+ let end = (start + 64).min(len);
+ for (hash, key) in
hashes[start..end].iter_mut().zip(&keys[start..end]) {
+ scatter_row(hash, key);
+ }
+ } else {
+ while word != 0 {
+ let i = start + word.trailing_zeros() as usize;
+ if i >= len {
+ break;
+ }
+ scatter_row(&mut hashes[i], &keys[i]);
+ word &= word - 1;
+ }
}
}
} else {
- for (hash, key) in
hashes_buffer.iter_mut().zip(array.keys().values()) {
- let idx = key.as_usize();
- if !HAS_NULL_VALUES || dict_values.is_valid(idx) {
- // SAFETY: `idx` is a valid Arrow dict key, in-bounds for
`dict_hashes`.
- let dict_hash = unsafe { *dict_hashes.get_unchecked(idx) };
- *hash = if MULTI_COL {
- combine_hashes(dict_hash, *hash)
- } else {
- dict_hash
- };
- }
+ for (hash, key) in hashes_buffer.iter_mut().zip(keys) {
+ scatter_row(hash, key);
}
}
}
```
<details><summary>Numbers, and one test to add with this loop</summary>
`create_hashes`, Int32 keys, 8,192 rows, 1,000 string values, 64 different
batches in rotation, median of 8 runs on an Apple M4 Pro, in µs:
| Nulls | 1 col: main / this PR / word loop | 3 cols: main / this PR / word
loop |
|---|---|---|
| 1 null | 5.72 / 5.51 / 2.65 | 18.0 / 18.6 / 10.3 |
| 0.1% | 5.96 / 5.54 / 2.99 | 18.1 / 18.7 / 10.6 |
| 1% | 6.71 / 5.88 / 4.55 | 19.8 / 19.7 / 14.9 |
| 3% | 7.63 / 6.04 / 5.65 | 22.9 / 20.1 / 18.3 |
| 50% | 28.7 / 4.86 / 4.98 | 87.8 / 13.5 / 13.4 |
| 97% | 7.01 / 2.65 / 2.71 | 20.9 / 5.86 / 6.14 |
This test panics if the `start >= len` check is removed, because an
all-valid word then starts past the end of a short buffer. Your current test
doesn't reach that shape:
```rs
#[test]
// Tests actual values of hashes, which are different if forcing collisions
#[cfg(not(feature = "force_hash_collisions"))]
fn create_hashes_for_dict_short_buffer_before_all_valid_word() {
// The word at row 64 is all valid and starts past the end of the buffer
let keys = Int32Array::from_iter((0..200).map(|row| (row !=
199).then_some(0)));
let values = Arc::new(StringArray::from(vec!["alpha"]));
let dict = DictionaryArray::try_new(keys, values).unwrap();
let random_state = RandomState::with_seed(0);
let mut hashes = vec![0; 200];
create_hashes([&dict as &dyn Array], &random_state, &mut
hashes).unwrap();
let mut short = vec![0; 10];
create_hashes([&dict as &dyn Array], &random_state, &mut short).unwrap();
assert_eq!(short, hashes[..10]);
}
```
</details>
--
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]
