Rich-T-kid commented on code in PR #12008:
URL: https://github.com/apache/arrow-rs/pull/12008#discussion_r4503271214
##########
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:
Changed the gate to use the actual remap count. Concat defaults to the
key-array length, interleave passes per-input lookup counts, and the new gate
test covers the mostly-null and repeated-selection cases. This is in bdd44b6963.
##########
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:
Raised the cutoff to 256. I also added contiguous and scattered benchmark
cases at 128, 256, 400, and 1024, and updated the description with the
crossover numbers. bdd44b6963
##########
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:
Rebased onto current main. The coalescer now reads through KeyMapping::get,
and merge_dictionary_values_with_masks uses the selected-index helper that
replaced get_masked_values. I kept the separate BatchCoalescer optimization out
of this PR and made that scope explicit in the description. bdd44b6963
##########
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:
Fixed this as suggested: assert_mapping now requires the Dense variant and
compares the whole Vec, so a missing or truncated entry cannot read back as
zero. bdd44b6963
--
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]
