okhsunrog commented on code in PR #10927:
URL: https://github.com/apache/arrow-rs/pull/10927#discussion_r3905604603
##########
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.
+ let dict = || {
+ let values: BinaryViewArray = (0..200u32)
+ .map(|i| Some(i.to_le_bytes().to_vec()))
+ .collect();
+ let keys = UInt8Array::from_iter_values(0..200);
+ DictionaryArray::<UInt8Type>::new(keys, Arc::new(values))
+ };
+
+ let combined = concat(&[&dict(), &dict()]).unwrap();
+ let combined = combined.as_dictionary::<UInt8Type>();
+
+ assert_eq!(combined.len(), 400);
+ assert_eq!(combined.values().data_type(), &DataType::BinaryView);
+ assert!(combined.values().len() < 400);
+
+ let values = combined.values().as_binary_view();
+ let actual: Vec<_> = combined
+ .keys()
+ .values()
+ .iter()
+ .map(|k| values.value(*k as usize))
+ .collect();
+ let expected: Vec<_> = (0..2)
+ .flat_map(|_| (0..200u32).map(|i| i.to_le_bytes().to_vec()))
+ .collect();
+ assert_eq!(actual, expected);
+ }
+
+ #[test]
+ fn concat_dictionary_merges_values_of_many_arrays() {
+ // Four dictionaries over the same 200 distinct values. Concatenating
+ // their values would need 800 keys, well past the u8 key range, so
this
+ // only succeeds if the merge deduplicates them.
+ let dicts: Vec<_> = (0..4)
+ .map(|d| {
+ let values: StringArray = (0..200).map(|i|
Some(format!("v{i}"))).collect();
+ let keys = UInt8Array::from_iter_values((0..200).map(|i| (i +
d * 17) as u8 % 200));
+ DictionaryArray::<UInt8Type>::new(keys, Arc::new(values))
+ })
+ .collect();
+ let refs: Vec<&dyn Array> = dicts.iter().map(|d| d as &dyn
Array).collect();
+
+ let combined = concat(&refs).unwrap();
+ let combined = combined.as_dictionary::<UInt8Type>();
+
+ assert_eq!(combined.len(), 800);
+ // Every key is addressable; how far below 200 the merge gets depends
on
+ // whether the best-effort interner sufficed or the exact retry ran
+ assert!(u8::try_from(combined.values().len()).is_ok());
Review Comment:
Right, this did read badly, the comment was wrong: the merged values can
never go below 200, there are exactly that many distinct ones. Four
dictionaries do overflow the interner, so the exact retry runs and the count
comes out at exactly 200. Asserting that now, with the comment rewritten.
--
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]
