Jefffrey commented on code in PR #11629:
URL: https://github.com/apache/arrow-rs/pull/11629#discussion_r4305576564
##########
arrow-select/src/interleave.rs:
##########
@@ -2229,6 +2240,137 @@ mod tests {
assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
}
+ fn dictionary_strings(array: &dyn Array) -> Vec<String> {
+ array
+ .as_dictionary::<Int8Type>()
+ .downcast_dict::<StringArray>()
+ .unwrap()
+ .into_iter()
+ .map(|v| v.unwrap().to_string())
+ .collect()
+ }
+
+ /// An Int8 dictionary with one key and 128 values, the most Int8 keys can
address.
+ /// Each call allocates its own values, so two dictionaries are not
pointer-equal.
+ fn int8_dictionary_128_values() -> ArrayRef {
+ let values = StringArray::from_iter_values((0..128).map(|i|
i.to_string()));
+ Arc::new(DictionaryArray::<Int8Type>::new(
+ Int8Array::from(vec![0]),
+ Arc::new(values),
+ ))
+ }
+
+ /// One list view row of size 1 per element of `values`.
+ fn single_element_list_views<O: OffsetSizeTrait>(
+ field: FieldRef,
+ values: ArrayRef,
+ ) -> GenericListViewArray<O> {
+ let len = values.len();
+ GenericListViewArray::new(
+ field,
+ (0..len).map(O::usize_as).collect::<Vec<_>>().into(),
+ vec![O::usize_as(1); len].into(),
+ values,
+ None,
+ )
+ }
+
+ fn test_interleave_list_view_dictionary_overflow<O: OffsetSizeTrait>(
+ indices: &[(usize, usize)],
+ ) {
+ // The two dictionaries hold 256 values together, more than Int8 keys
can
+ // address, but the selected rows only reference one distinct value.
+ let field = Arc::new(Field::new_list_field(
+ DataType::Dictionary(Box::new(DataType::Int8),
Box::new(DataType::Utf8)),
+ true,
+ ));
+ let a = single_element_list_views::<O>(field.clone(),
int8_dictionary_128_values());
+ let b = single_element_list_views::<O>(field,
int8_dictionary_128_values());
+
+ let result = interleave(&[&a, &b], indices).unwrap();
+
+ let result = result.as_list_view::<O>();
+ assert_eq!(result.len(), indices.len());
+ for i in 0..result.len() {
+ assert_eq!(dictionary_strings(result.value(i).as_ref()), ["0"]);
+ }
+ }
+
+ #[test]
+ fn test_list_view_dictionary_overflow_copy() {
+ // The selected rows are cheaper to copy than to concat the backing
arrays
+ test_interleave_list_view_dictionary_overflow::<i32>(&[(0, 0), (1,
0)]);
Review Comment:
we can simplify the tests to just test i32 offset case; largelistview takes
exact same codepath anyway
--
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]
