Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix merge_dictionary_values in selection kernels #4833

Merged
merged 2 commits into from
Sep 19, 2023
Merged

Fix merge_dictionary_values in selection kernels #4833

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2372928
Fix merge_dictionary_values in selection kernels
tustvold Sep 18, 2023
fee2087
Merge remote-tracking branch 'upstream/master' into fix-interleave-dict
tustvold Sep 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 23 additions & 9 deletions arrow-select/src/dictionary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ pub fn merge_dictionary_values<K: ArrowDictionaryKeyType>(
) -> Result<MergedDictionaries<K>, ArrowError> {
let mut num_values = 0;

let mut values = Vec::with_capacity(dictionaries.len());
let mut values_arrays = Vec::with_capacity(dictionaries.len());
let mut value_slices = Vec::with_capacity(dictionaries.len());

for (idx, dictionary) in dictionaries.iter().enumerate() {
Expand All @@ -164,11 +164,13 @@ pub fn merge_dictionary_values<K: ArrowDictionaryKeyType>(
(None, None) => None,
};
let keys = dictionary.keys().values();
let values_mask = compute_values_mask(keys, key_mask.as_ref());
let v = dictionary.values().as_ref();
num_values += v.len();
value_slices.push(get_masked_values(v, &values_mask));
values.push(v)
let values = dictionary.values().as_ref();
let values_mask = compute_values_mask(keys, key_mask.as_ref(), values.len());

let masked_values = get_masked_values(values, &values_mask);
num_values += masked_values.len();
value_slices.push(masked_values);
values_arrays.push(values)
}

// Map from value to new index
Expand Down Expand Up @@ -202,7 +204,7 @@ pub fn merge_dictionary_values<K: ArrowDictionaryKeyType>(

Ok(MergedDictionaries {
key_mappings,
values: interleave(&values, &indices)?,
values: interleave(&values_arrays, &indices)?,
})
}

Expand All @@ -211,9 +213,10 @@ pub fn merge_dictionary_values<K: ArrowDictionaryKeyType>(
fn compute_values_mask<K: ArrowNativeType>(
keys: &ScalarBuffer<K>,
mask: Option<&BooleanBuffer>,
max_key: usize,
) -> BooleanBuffer {
let mut builder = BooleanBufferBuilder::new(keys.len());
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the bug, the length of the values_mask should match that of the values array not the keys. In all of the existing tests the number of keys outnumbered the number of values, not helped by the heuristic to trigger merging being based off the total number of keys outnumbering the total number of values.

builder.advance(keys.len());
let mut builder = BooleanBufferBuilder::new(max_key);
builder.advance(max_key);

match mask {
Some(n) => n
Expand Down Expand Up @@ -330,4 +333,15 @@ mod tests {
assert_eq!(&merged.key_mappings[0], &[0, 0, 0, 1, 0]);
assert_eq!(&merged.key_mappings[1], &[]);
}

#[test]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I verified that without the changes in this PR this test fails like this:

thread 'dictionary::tests::test_merge_keys_smaller' panicked at 'assertion failed: `(left == right)`
  left: `StringArray
[
]`,
 right: `StringArray
[
  "b",
]`', arrow-select/src/dictionary.rs:342:9

fn test_merge_keys_smaller() {
let values = StringArray::from_iter_values(["a", "b"]);
let keys = Int32Array::from_iter_values([1]);
let a = DictionaryArray::new(keys, Arc::new(values));

let merged = merge_dictionary_values(&[&a], None).unwrap();
let expected = StringArray::from(vec!["b"]);
assert_eq!(merged.values.as_ref(), &expected);
}
}