jayzhan211 commented on code in PR #12008:
URL: https://github.com/apache/arrow-rs/pull/12008#discussion_r4504315587
##########
arrow-select/src/coalesce/dictionary.rs:
##########
@@ -153,12 +153,18 @@ impl<K: ArrowDictionaryKeyType>
InProgressDictionaryArray<K> {
match &chunk.nulls {
Some(nulls) =>
keys.extend(chunk.keys.iter().enumerate().map(|(offset, key)| {
if nulls.is_valid(offset) {
- mapping[key.as_usize()]
+ mapping
Review Comment:
Optional cleanup: `merge_dictionary_values_with_masks` only returns dense
mappings, so matching once per chunk keeps this loop the same as on main
instead of matching the enum and calling `expect` for every key. The
`coalesce_kernels` dictionary benchmarks were within noise either way, so this
is only about keeping the hot loop simple.
```diff
-use crate::dictionary::{child_values_ptr_eq,
merge_dictionary_values_with_masks};
+use crate::dictionary::{KeyMapping, child_values_ptr_eq,
merge_dictionary_values_with_masks};
@@
- let mapping = &merged.key_mappings[chunk.source];
+ let KeyMapping::Dense(mapping) =
&merged.key_mappings[chunk.source] else {
+ unreachable!("merge_dictionary_values_with_masks returns
dense mappings");
+ };
match &chunk.nulls {
Some(nulls) =>
keys.extend(chunk.keys.iter().enumerate().map(|(offset, key)| {
if nulls.is_valid(offset) {
- mapping
- .get(key.as_usize())
- .expect("referenced dictionary key should be
mapped")
+ mapping[key.as_usize()]
} else {
K::Native::default()
}
})),
- None => keys.extend(chunk.keys.iter().map(|key| {
- mapping
- .get(key.as_usize())
- .expect("referenced dictionary key should be
mapped")
- })),
+ None => keys.extend(chunk.keys.iter().map(|key|
mapping[key.as_usize()])),
}
```
##########
arrow-select/src/concat.rs:
##########
@@ -1567,6 +1577,41 @@ mod tests {
assert!((30..40).contains(&values_len), "{values_len}")
}
+ #[test]
+ fn test_concat_sparse_primitive_dictionaries_with_nulls() {
+ let values_len = 3 *
crate::dictionary::MAX_DENSE_MAPPING_VALUES_PER_ROW + 1;
+ let last = values_len as i32 - 1;
+ let make_dictionary = |offset, null_value, keys| {
+ let values = Int32Array::from_iter(
+ (0..values_len).map(|idx| (idx != null_value).then_some(offset
+ idx as i32)),
+ );
+ DictionaryArray::<Int32Type>::new(Int32Array::from(keys),
Arc::new(values))
+ };
+ let a = make_dictionary(0, 1, vec![Some(last), None, Some(1)]);
+ let b = make_dictionary(10_000, 2, vec![Some(2), Some(last - 1),
None]);
Review Comment:
Could you slice the inputs in this test? Since
`test_concat_oversized_dictionary_slices` was replaced, no test sends a sliced
dictionary through the sparse path, which is the case in this PR's title. With
the change below both inputs, including their null buffers, start at offset 1,
and the test also checks that the gate counts the sliced length: the unsliced
arrays have 4 keys, and 4 × 256 ≥ 769 would keep them dense.
```suggestion
// Slicing moves the keys and their nulls to a non-zero offset; the
unsliced
// arrays have 4 keys each and would take the dense path.
let a = make_dictionary(0, 1, vec![Some(0), Some(last), None,
Some(1)]).slice(1, 3);
let b =
make_dictionary(10_000, 2, vec![Some(0), Some(2), Some(last -
1), None]).slice(1, 3);
```
--
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]
