Rich-T-kid commented on code in PR #10927:
URL: https://github.com/apache/arrow-rs/pull/10927#discussion_r3925338378
##########
arrow-select/src/dictionary.rs:
##########
@@ -257,13 +266,84 @@ pub(crate) fn merge_dictionary_values<K:
ArrowDictionaryKeyType>(
values_arrays.push(values)
}
- // Map from value to new index
- let mut interner = Interner::new(num_values);
// Interleave indices for new values array
let mut indices = Vec::with_capacity(num_values);
- // Compute the mapping for each dictionary
- let key_mappings = dictionaries
+ // Map from value to new index, best-effort: a hash collision evicts the
+ // previous occupant, so the same value may be assigned more than one key
+ let mut interner = Interner::new(num_values);
+ let interned = compute_key_mappings(
+ dictionaries,
+ &value_slices,
+ |dictionary_idx, value_idx, value| {
+ interner
+ .intern(value, || match K::Native::from_usize(indices.len()) {
+ Some(idx) => {
+ indices.push((dictionary_idx, value_idx));
+ Ok(idx)
+ }
+ None => Err(ArrowError::DictionaryKeyOverflowError),
+ })
+ .copied()
+ },
+ );
+
+ let key_mappings = match interned {
+ Ok(key_mappings) => key_mappings,
+ // The duplicates left behind by the interner's hash collisions can
push
+ // the output past what the key type can address even though the
distinct
+ // values would have fit. Retry with exact deduplication, which
allocates
+ // exactly one key per distinct value at the cost of a hash map.
Review Comment:
a bit unrelated to this PR but I think it'd be useful to add some of this
info to the `Internern::intern()`. its easier to understand why we retry on
DictionaryKeyOverFlowError if the `Internern::intern()` was a bit more clear.
##########
arrow-select/src/concat.rs:
##########
@@ -1708,6 +1708,124 @@ mod tests {
assert_eq!(array.logical_null_count(), 10);
}
+ #[test]
+ fn concat_string_view_dictionary_merges_duplicate_values() {
+ // Two independently-built `Dictionary<UInt8, Utf8View>` arrays
holding the
+ // same 200 distinct values. Naively concatenating their dictionaries
yields
+ // 400 entries, which overflows the u8 key range, but the distinct
values do
+ // fit -- so the values must be merged and deduplicated instead. This
mirrors
+ // a `Dictionary<UInt16, Utf8View>` column read in several partitions,
each
+ // building its own dictionary, and then combined.
+ let dict = |offset: usize| {
+ let values: StringViewArray = (0..200).map(|i|
Some(format!("v{i}"))).collect();
+ let keys = UInt8Array::from_iter_values((0..200).map(|i| (i +
offset) as u8 % 200));
+ DictionaryArray::<UInt8Type>::new(keys, Arc::new(values))
+ };
+ let (a, b) = (dict(0), dict(7));
+
+ let combined = concat(&[&a, &b]).unwrap();
+ let combined = combined.as_dictionary::<UInt8Type>();
+
+ assert_eq!(combined.len(), 400);
+ assert_eq!(combined.values().data_type(), &DataType::Utf8View);
+ assert!(combined.values().len() < 400);
+
+ let values = combined.values().as_string_view();
+ let actual: Vec<_> = combined
+ .keys()
+ .values()
+ .iter()
+ .map(|k| values.value(*k as usize))
+ .collect();
+ let expected: Vec<_> = [&a, &b]
+ .iter()
+ .flat_map(|d| {
+ let v = d.values().as_string_view();
+ d.keys()
+ .values()
+ .iter()
+ .map(|k| v.value(*k as usize))
+ .collect::<Vec<_>>()
+ })
+ .collect();
+ assert_eq!(actual, expected);
+ }
+
+ #[test]
+ fn concat_binary_view_dictionary_merges_duplicate_values() {
+ // Same as `concat_string_view_dictionary_merges_duplicate_values`,
for the
+ // other view-typed dictionary value layout.
Review Comment:
```rust
/// Downcast this to a [`StringViewArray`]
///
/// # Panics
///
/// Panics if this is not a [`StringViewArray`]
fn as_string_view(&self) -> &StringViewArray {
self.as_byte_view_opt().expect("string view array")
}
/// Downcast this to a [`BinaryViewArray`]
///
/// # Panics
///
/// Panics if this is not a [`BinaryViewArray`]
fn as_binary_view(&self) -> &BinaryViewArray {
self.as_byte_view_opt().expect("binary view array")
}
```
these just perform a a downcast so that `masked_byte_views()` receives a
`GenericByteViewArray` the code paths are identical
https://github.com/apache/arrow-rs/blob/027b45f2a982832c93d159de6718d69d0e8b60e3/arrow-array/src/array/byte_view_array.rs#L1125
https://github.com/apache/arrow-rs/blob/027b45f2a982832c93d159de6718d69d0e8b60e3/arrow-array/src/array/byte_view_array.rs#L1166
this isn't too important but i'm in favor or removing it. would be nice to
get a third opinion
--
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]
