mbutrovich commented on code in PR #11407:
URL: https://github.com/apache/arrow-rs/pull/11407#discussion_r4275023113
##########
arrow-cmp/src/lib.rs:
##########
@@ -798,6 +867,186 @@ mod tests {
assert_eq!(Ordering::Greater, cmp(1, 3));
}
+ #[test]
+ fn dict_fixed_size_binary_fast_path_derefs_values_both_directions() {
+ let shared_keys = || Int32Array::from(vec![0_i32, 1]);
+ let left_values =
+ FixedSizeBinaryArray::try_from_iter([b"zz".as_ref(),
b"aa".as_ref()].into_iter())
+ .unwrap();
+ let right_values =
+ FixedSizeBinaryArray::try_from_iter([b"aa".as_ref(),
b"zz".as_ref()].into_iter())
+ .unwrap();
+ let left = DictionaryArray::<Int32Type>::new(shared_keys(),
Arc::new(left_values));
+ let right = DictionaryArray::<Int32Type>::new(shared_keys(),
Arc::new(right_values));
+
+ let ascending = make_comparator(&left, &right,
SortOptions::default()).unwrap();
+ assert_eq!(ascending(0, 0), Ordering::Greater);
+ assert_eq!(ascending(1, 1), Ordering::Less);
+
+ let descending = make_comparator(
+ &left,
+ &right,
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )
+ .unwrap();
+ assert_eq!(descending(0, 0), Ordering::Less);
+ assert_eq!(descending(1, 1), Ordering::Greater);
+ }
+
+ #[test]
+ fn dict_bytes_fast_path_derefs_values_both_directions() {
+ let shared_keys = || Int32Array::from(vec![0_i32, 1]);
+ let left = DictionaryArray::<Int32Type>::new(
+ shared_keys(),
+ Arc::new(StringArray::from(vec!["z", "a"])),
+ );
+ let right = DictionaryArray::<Int32Type>::new(
+ shared_keys(),
+ Arc::new(StringArray::from(vec!["a", "z"])),
+ );
+
+ let ascending = make_comparator(&left, &right,
SortOptions::default()).unwrap();
+ assert_eq!(ascending(0, 0), Ordering::Greater);
+ assert_eq!(ascending(1, 1), Ordering::Less);
+
+ let descending = make_comparator(
+ &left,
+ &right,
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )
+ .unwrap();
+ assert_eq!(descending(0, 0), Ordering::Less);
+ assert_eq!(descending(1, 1), Ordering::Greater);
+ }
+
+ #[test]
+ fn dict_bytes_fallback_handles_key_and_value_nulls() {
Review Comment:
Should we drop `fallback` from this name, and from the ones on lines 948 and
1002? With the separate no-null closures gone, there's no fallback path left in
`compare_dict_bytes` or `compare_dict_fixed_size_binary`. A reader could take
"fallback" to mean the generic `make_comparator` path in `compare_dict`, and
these tests don't reach that path.
```suggestion
fn dict_bytes_handles_key_and_value_nulls() {
```
##########
arrow-cmp/src/lib.rs:
##########
@@ -220,6 +220,25 @@ fn compare_dict<K: ArrowDictionaryKeyType>(
) -> Result<DynComparator, ArrowError> {
let left = left.as_dictionary::<K>();
let right = right.as_dictionary::<K>();
+ let value_dtype = left.values().data_type();
+ if value_dtype == right.values().data_type() {
+ match value_dtype {
+ DataType::Utf8 => return Ok(compare_dict_bytes::<K,
Utf8Type>(left, right, opts)),
+ DataType::LargeUtf8 => {
+ return Ok(compare_dict_bytes::<K, LargeUtf8Type>(left, right,
opts));
+ }
+ DataType::Binary => {
+ return Ok(compare_dict_bytes::<K, BinaryType>(left, right,
opts));
+ }
+ DataType::LargeBinary => {
+ return Ok(compare_dict_bytes::<K, LargeBinaryType>(left,
right, opts));
+ }
Review Comment:
Could we add a test that goes through the `LargeUtf8`, `Binary` and
`LargeBinary` arms? The new tests all build `StringArray` or
`FixedSizeBinaryArray` values, so these three arms still aren't reached by any
test (this was part of
https://github.com/apache/arrow-rs/pull/11407#discussion_r4264007280). A
generic helper in the style of
[`test_bytes_impl`](https://github.com/apache/arrow-rs/blob/b76b9b62f6015d74e8b4b70ca4c647a9082dd33a/arrow-cmp/src/lib.rs#L1254-L1270)
would cover all four byte types in a few lines:
```rust
fn test_dict_bytes_impl<T: ByteArrayType>() {
let offsets = OffsetBuffer::from_lengths([1, 1]);
let values = GenericByteArray::<T>::new(offsets, b"za".into(), None);
let keys = Int32Array::from(vec![Some(0), None, Some(1)]);
let dict = DictionaryArray::<Int32Type>::new(keys, Arc::new(values));
let cmp = make_comparator(&dict, &dict, SortOptions::default()).unwrap();
assert_eq!(cmp(0, 2), Ordering::Greater);
assert_eq!(cmp(1, 2), Ordering::Less);
assert_eq!(cmp(1, 1), Ordering::Equal);
}
#[test]
fn test_dict_bytes() {
test_dict_bytes_impl::<Utf8Type>();
test_dict_bytes_impl::<LargeUtf8Type>();
test_dict_bytes_impl::<BinaryType>();
test_dict_bytes_impl::<LargeBinaryType>();
}
```
The out-of-bounds test also only covers `compare_dict_bytes`. What do you
think about adding the same `#[should_panic]` case for a `FixedSizeBinary`
dictionary?
--
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]
