jayzhan211 commented on code in PR #12008:
URL: https://github.com/apache/arrow-rs/pull/12008#discussion_r4500232766
##########
arrow-select/src/dictionary.rs:
##########
@@ -140,9 +140,47 @@ impl<'a, V> Interner<'a, V> {
}
}
+// Sparse hashing pays off only when retained values greatly outnumber
selected rows.
+const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 128;
Review Comment:
With scattered keys, concat is 18-23% slower just above 128:1: 12 x 682 rows
at 129:1 goes from 216 µs to 256 µs, and 2 x 6,900 rows goes from 412 µs to 487
µs. The slice benchmark and the #11885 repro use contiguous keys, where sparse
wins from 129:1, so scattered keys cross over later, near 256:1. Could we use
256? That gives up as much as 23% on contiguous keys between 129:1 and 256:1,
and the DataFusion shape in #11885 is about 880:1, so it isn't affected.
```suggestion
const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 256;
```
<details>
<summary>Head / base ratio by values per selected row and key order</summary>
Median of 5 rounds of best-of-9 on an M4 Pro, spread mostly under 12%. All
keys valid. Each cell is contiguous / scattered.
| Values per row | 12 x 682 rows | 3 x 2,730 rows | 2 x 6,900 rows |
|---:|---:|---:|---:|
| 128 (dense) | 0.97 / 1.02 | 0.96 / 0.98 | 0.95 / 0.97 |
| 129 | 0.97 / 1.19 | 0.92 / 1.23 | 0.93 / 1.18 |
| 200 | 0.85 / 1.01 | 0.82 / 1.10 | 0.78 / 1.09 |
| 256 | 0.77 / 0.95 | 0.79 / 1.05 | 0.80 / 0.96 |
| 400 | 0.64 / 0.76 | 0.57 / 0.75 | 0.63 / 0.84 |
| 1024 | 0.40 / 0.53 | 0.35 / 0.50 | 0.44 / 0.56 |
</details>
##########
arrow-select/src/dictionary.rs:
##########
@@ -249,11 +287,23 @@ pub(crate) fn merge_dictionary_values<K:
ArrowDictionaryKeyType>(
};
let keys = dictionary.keys().values();
let values = dictionary.values().as_ref();
- let values_mask = compute_values_mask(keys, key_mask, values.len());
-
- let masked_values = get_masked_values(values, &values_mask);
- num_values += masked_values.len();
- value_slices.push(masked_values);
+ let selected_rows = key_mask
+ .map(BooleanBuffer::count_set_bits)
+ .unwrap_or(keys.len());
+ let use_sparse_mapping =
+ values.len() >
selected_rows.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW);
Review Comment:
The sparse check counts selected rows, but the remap does one hash lookup
per key, so mostly-null and repeated inputs get slower: concat of 2 x 8,192
rows with 100k values and 500 valid goes from 39 µs to 57 µs, and 2 x 1M rows
with 1M values and 7,000 valid goes from 2.6 ms to 28.9 ms. Could the check
count key lookups instead? With the change below those two take 38 µs and 2.4
ms, and the 1M-value repro stays at 170 µs.
```diff
--- a/arrow-select/src/dictionary.rs
+++ b/arrow-select/src/dictionary.rs
masks: Option<&[BooleanBuffer]>,
+ key_lookups: Option<&[usize]>,
) -> Result<MergedDictionaries<K>, ArrowError> {
- let selected_rows = key_mask
- .map(BooleanBuffer::count_set_bits)
- .unwrap_or(keys.len());
+ // The caller remaps once per lookup, which can far exceed the
selected rows.
+ let lookups = key_lookups.map_or(keys.len(), |lookups|
lookups[idx]);
let use_sparse_mapping =
- values.len() >
selected_rows.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW);
+ values.len() >
lookups.saturating_mul(MAX_DENSE_MAPPING_VALUES_PER_ROW);
--- a/arrow-select/src/interleave.rs
+++ b/arrow-select/src/interleave.rs
- let masks: Vec<_> = dictionaries
+ let (masks, key_lookups): (Vec<_>, Vec<_>) = dictionaries
+ let mut lookups = 0;
for (_, key_idx) in indices.iter().filter(|(a, _)| *a == a_idx)
{
key_mask.set_bit(*key_idx, true);
+ lookups += 1;
}
- key_mask.finish()
+ (key_mask.finish(), lookups)
})
- .collect();
+ .unzip();
- let merged = merge_dictionary_values(&dictionaries, Some(&masks))?;
+ let merged = merge_dictionary_values(&dictionaries, Some(&masks),
Some(&key_lookups))?;
--- a/arrow-select/src/concat.rs
+++ b/arrow-select/src/concat.rs
- let merged = merge_dictionary_values(&dictionaries, None)?;
+ let merged = merge_dictionary_values(&dictionaries, None, None)?;
```
<details>
<summary>Timings per input shape, and a test that separates lookups from
selected rows</summary>
Median of 3 rounds, each the best of at least 9 runs, on an M4 Pro under
load, so ignore differences under 10%. Null-slot keys are in range.
| Input | Base | This PR | With the change |
|---|---:|---:|---:|
| concat 12 x 682 rows, 1M values | 502 µs | 164 µs | 170 µs |
| concat 2 x 8,192 rows, 100k values, 500 valid | 39 µs | 57 µs | 38 µs |
| concat 2 x 1M rows, 1M values, 7,000 valid | 2.6 ms | 28.9 ms | 2.4 ms |
| concat 2 x 1M rows, 1M values, 10 valid | 1.7 ms | 3.7 ms | 1.5 ms |
| interleave 1M indices over 100 rows, 2 x 1M values | 2.7 ms | 4.0 ms | 2.8
ms |
| interleave 8,192 indices over 100 rows, 2 x 1M values | 104 µs | 58 µs |
106 µs |
The last row is what the change gives up: heavily repeated rows go back to
the dense path.
```rs
#[test]
fn test_sparse_gate_counts_key_lookups() {
let values: ArrayRef = Arc::new(StringArray::from_iter_values(
(0..8 * MAX_DENSE_MAPPING_VALUES_PER_ROW).map(|idx| idx.to_string()),
));
let is_sparse = |dictionary: &DictionaryArray<Int32Type>,
masks: Option<&[BooleanBuffer]>,
lookups: Option<&[usize]>| {
let merged = merge_dictionary_values(&[dictionary], masks,
lookups).unwrap();
matches!(merged.key_mappings[0], KeyMapping::Sparse(_))
};
// 1 valid row of 10: over the gate for 1 row, but concat looks up all
10 rows.
let keys = Int32Array::from_iter((0..10).map(|idx| (idx ==
0).then_some(999)));
let mostly_null = DictionaryArray::new(keys, Arc::clone(&values));
assert!(!is_sparse(&mostly_null, None, None));
// 7 lookups stay under the gate, keeping the sparse path.
let keys = Int32Array::from_iter_values([999; 7]);
let short = DictionaryArray::new(keys, Arc::clone(&values));
assert!(is_sparse(&short, None, None));
// Interleave selects 1 row but may look it up repeatedly.
let mask = BooleanBuffer::collect_bool(7, |idx| idx == 0);
assert!(!is_sparse(
&short,
Some(std::slice::from_ref(&mask)),
Some(&[10])
));
assert!(is_sparse(&short, Some(&[mask]), Some(&[7])));
}
```
</details>
##########
arrow-select/src/dictionary.rs:
##########
@@ -140,9 +140,47 @@ impl<'a, V> Interner<'a, V> {
}
}
+// Sparse hashing pays off only when retained values greatly outnumber
selected rows.
+const MAX_DENSE_MAPPING_VALUES_PER_ROW: usize = 128;
+
+pub(crate) enum KeyMapping<K: ArrowDictionaryKeyType> {
+ Dense(Vec<K::Native>),
+ Sparse(HashMap<usize, K::Native, RandomState>),
+}
+
+impl<K: ArrowDictionaryKeyType> KeyMapping<K> {
+ fn new_dense(len: usize) -> Self {
+ Self::Dense(vec![K::Native::ZERO; len])
+ }
+
+ fn new_sparse(capacity: usize) -> Self {
+ Self::Sparse(HashMap::with_capacity_and_hasher(
+ capacity,
+ RandomState::with_seeds(0, 0, 0, 0),
+ ))
+ }
+
+ fn insert(&mut self, old_key: usize, new_key: K::Native) {
+ match self {
+ Self::Dense(mapping) => mapping[old_key] = new_key,
+ Self::Sparse(mapping) => {
+ mapping.insert(old_key, new_key);
+ }
+ }
+ }
+
+ #[inline]
+ pub fn get(&self, old_key: usize) -> Option<K::Native> {
+ match self {
+ Self::Dense(mapping) => mapping.get(old_key).copied(),
+ Self::Sparse(mapping) => mapping.get(&old_key).copied(),
+ }
+ }
+}
+
pub(crate) struct MergedDictionaries<K: ArrowDictionaryKeyType> {
- /// Provides `key_mappings[`array_idx`][`old_key`] -> new_key`
- pub key_mappings: Vec<Vec<K::Native>>,
+ /// Maps each input dictionary key to its merged key.
+ pub key_mappings: Vec<KeyMapping<K>>,
Review Comment:
This needs more than a text-conflict rebase: since #11743,
`coalesce/dictionary.rs` indexes `mapping[key.as_usize()]` and
`merge_dictionary_values_with_masks` calls `get_masked_values`, and this PR
changes the mapping's type and deletes that helper. Could you rebase and re-run
the ClickBench comparison? Those numbers pinned `258b0d2031`, before
`BatchCoalescer` stopped calling `concat`, so the Q16-Q34 gains may not carry
over.
##########
arrow-select/src/concat.rs:
##########
@@ -1559,6 +1569,35 @@ mod tests {
assert!((30..40).contains(&values_len), "{values_len}")
}
+ #[test]
+ fn test_concat_oversized_dictionary_slices() {
Review Comment:
This test and `test_interleave_oversized_dictionaries` hardcode 1,000 values
and never assert the sparse path ran, so raising the threshold can move them to
the dense path with no failure. Could they derive sizes from
`MAX_DENSE_MAPPING_VALUES_PER_ROW`, and add sparse cases for primitive values,
null dictionary values, and null keys through both `concat` and `interleave`? A
randomized base-vs-head run over those shapes found no differences in 60,000
operations, so this is only about keeping them covered.
##########
arrow-select/src/dictionary.rs:
##########
@@ -374,12 +448,18 @@ mod tests {
use super::*;
use arrow_array::cast::as_string_array;
- use arrow_array::types::Int8Type;
- use arrow_array::types::Int32Type;
- use arrow_array::{DictionaryArray, Int8Array, Int32Array, StringArray};
+ use arrow_array::types::{Int8Type, Int32Type, UInt8Type};
+ use arrow_array::{DictionaryArray, Int8Array, Int32Array, StringArray,
UInt8Array};
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer};
use std::sync::Arc;
+ fn assert_mapping(mapping: &KeyMapping<Int32Type>, expected: &[i32]) {
+ let actual = (0..expected.len())
+ .map(|idx| mapping.get(idx).unwrap_or_default())
+ .collect::<Vec<_>>();
+ assert_eq!(actual, expected);
+ }
Review Comment:
`assert_mapping` no longer checks length or variant and reads a missing
entry as 0, so the `&[]` call at line 595 can't fail and these tests would
still pass if an input moved to the sparse path. Comparing the whole dense
`Vec` keeps the old assertions:
```suggestion
fn assert_mapping(mapping: &KeyMapping<Int32Type>, expected: &[i32]) {
let KeyMapping::Dense(actual) = mapping else {
panic!("expected a dense mapping");
};
assert_eq!(actual, expected);
}
```
--
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]
