mbutrovich commented on code in PR #11407:
URL: https://github.com/apache/arrow-rs/pull/11407#discussion_r4264003462
##########
arrow-cmp/src/lib.rs:
##########
@@ -234,6 +253,93 @@ fn compare_dict<K: ArrowDictionaryKeyType>(
Ok(f)
}
+#[inline]
+fn compare_dict_bytes<K: ArrowDictionaryKeyType, T: ByteArrayType>(
+ left: &DictionaryArray<K>,
+ right: &DictionaryArray<K>,
+ opts: SortOptions,
+) -> DynComparator {
+ let left_values = left.values().as_bytes::<T>().clone();
+ let right_values = right.values().as_bytes::<T>().clone();
+ let left_keys = left.keys().values().clone();
+ let right_keys = right.keys().values().clone();
+ // SAFETY: row indices are in-bounds per DynComparator contract; keys are
+ // in-bounds for values per DictionaryArray::try_new.
Review Comment:
Is there a `DynComparator` contract that says row indices are in bounds? I
couldn't find one.
[`DynComparator`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L71-L72)
and
[`make_comparator`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L464-L470)
only say "arbitrary indices", and `make_comparator` is a safe function that
returns a safe closure. As I read it, `get_unchecked(left_row)` lets safe code
read past the end of the keys buffer.
On the head commit, `cmp(2, 0)` on a two-row `Dictionary(Int32, Utf8)`
returns `Ordering::Equal` in a release build. In a debug build it stops at
std's `unsafe precondition(s) violated: slice::get_unchecked requires that the
index is within the slice`. On `main` the same call panics at
[`left_keys[i]`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L230).
The crate already treats an out-of-bounds index as a panic.
[`compare_byte_view_values`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L197-L204)
asserts before its unchecked access, and
[`test_union_out_of_bounds`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L1643-L1644)
expects `index out of bounds`.
Should we index the keys with `left_keys[left_row]` and keep
`value_unchecked` only for the key-to-value lookup? The `DictionaryArray` key
invariant covers that lookup, and `compare` only calls the closure for rows
that are valid on both sides. The checks add a load, a shift, a compare and a
branch per side, ahead of the `memcmp` call:
```
ldr x8, [x0, #256]
lsr x8, x8, #2 ; left_keys.len()
cmp x1, x8
b.hs LBB274_3 ; panic_bounds_check
```
##########
arrow-cmp/src/lib.rs:
##########
@@ -234,6 +253,93 @@ fn compare_dict<K: ArrowDictionaryKeyType>(
Ok(f)
}
+#[inline]
+fn compare_dict_bytes<K: ArrowDictionaryKeyType, T: ByteArrayType>(
+ left: &DictionaryArray<K>,
+ right: &DictionaryArray<K>,
+ opts: SortOptions,
+) -> DynComparator {
+ let left_values = left.values().as_bytes::<T>().clone();
+ let right_values = right.values().as_bytes::<T>().clone();
+ let left_keys = left.keys().values().clone();
+ let right_keys = right.keys().values().clone();
+ // SAFETY: row indices are in-bounds per DynComparator contract; keys are
+ // in-bounds for values per DictionaryArray::try_new.
+ if left.logical_null_count() == 0 && right.logical_null_count() == 0 {
+ return if opts.descending {
+ Box::new(move |left_row, right_row| unsafe {
+ let left_bytes: &[u8] = left_values
+
.value_unchecked(left_keys.get_unchecked(left_row).as_usize())
+ .as_ref();
+ let right_bytes: &[u8] = right_values
+
.value_unchecked(right_keys.get_unchecked(right_row).as_usize())
+ .as_ref();
+ right_bytes.cmp(left_bytes)
+ })
+ } else {
+ Box::new(move |left_row, right_row| unsafe {
+ let left_bytes: &[u8] = left_values
+
.value_unchecked(left_keys.get_unchecked(left_row).as_usize())
+ .as_ref();
+ let right_bytes: &[u8] = right_values
+
.value_unchecked(right_keys.get_unchecked(right_row).as_usize())
+ .as_ref();
+ left_bytes.cmp(right_bytes)
+ })
+ };
+ }
+ // SAFETY: same invariants as above.
+ compare(left, right, opts, move |left_row, right_row| unsafe {
+ let left_bytes: &[u8] = left_values
+ .value_unchecked(left_keys.get_unchecked(left_row).as_usize())
+ .as_ref();
+ let right_bytes: &[u8] = right_values
+ .value_unchecked(right_keys.get_unchecked(right_row).as_usize())
+ .as_ref();
+ left_bytes.cmp(right_bytes)
+ })
+}
+
+#[inline]
+fn compare_dict_fixed_size_binary<K: ArrowDictionaryKeyType>(
+ left: &DictionaryArray<K>,
+ right: &DictionaryArray<K>,
+ opts: SortOptions,
+) -> DynComparator {
+ let left_values = left.values().as_fixed_size_binary().clone();
+ let right_values = right.values().as_fixed_size_binary().clone();
+ let left_keys = left.keys().values().clone();
+ let right_keys = right.keys().values().clone();
+ // SAFETY: row indices are in-bounds per DynComparator contract; keys are
+ // in-bounds for values per DictionaryArray::try_new.
+ if left.logical_null_count() == 0 && right.logical_null_count() == 0 {
+ return if opts.descending {
+ Box::new(move |left_row, right_row| unsafe {
+ let left_bytes =
+
left_values.value_unchecked(left_keys.get_unchecked(left_row).as_usize());
+ let right_bytes =
+
right_values.value_unchecked(right_keys.get_unchecked(right_row).as_usize());
+ right_bytes.cmp(left_bytes)
+ })
+ } else {
+ Box::new(move |left_row, right_row| unsafe {
+ let left_bytes =
+
left_values.value_unchecked(left_keys.get_unchecked(left_row).as_usize());
+ let right_bytes =
+
right_values.value_unchecked(right_keys.get_unchecked(right_row).as_usize());
+ left_bytes.cmp(right_bytes)
+ })
+ };
+ }
+ // SAFETY: same invariants as above.
+ compare(left, right, opts, move |left_row, right_row| unsafe {
+ let left_bytes =
left_values.value_unchecked(left_keys.get_unchecked(left_row).as_usize());
+ let right_bytes =
+
right_values.value_unchecked(right_keys.get_unchecked(right_row).as_usize());
+ left_bytes.cmp(right_bytes)
+ })
+}
Review Comment:
This applies the same two changes to the `FixedSizeBinary` path.
```suggestion
fn compare_dict_fixed_size_binary<K: ArrowDictionaryKeyType>(
left: &DictionaryArray<K>,
right: &DictionaryArray<K>,
opts: SortOptions,
) -> DynComparator {
let left_values = left.values().as_fixed_size_binary().clone();
let right_values = right.values().as_fixed_size_binary().clone();
let left_keys = left.keys().values().clone();
let right_keys = right.keys().values().clone();
compare(left, right, opts, move |left_row, right_row| {
let left_key = left_keys[left_row].as_usize();
let right_key = right_keys[right_row].as_usize();
debug_assert!(left_key < left_values.len() && right_key <
right_values.len());
// SAFETY: `compare` calls this only for rows that are valid on both
sides, and
// `DictionaryArray` guarantees that every valid key indexes into
its values.
let (left_bytes, right_bytes) = unsafe {
(
left_values.value_unchecked(left_key),
right_values.value_unchecked(right_key),
)
};
left_bytes.cmp(right_bytes)
})
}
```
##########
arrow-cmp/src/lib.rs:
##########
@@ -798,6 +904,83 @@ 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() {
+ // Left carries a logical null via a null KEY at row 1.
+ let left_keys = Int32Array::from(vec![Some(0_i32), None, Some(1)]);
+ let left_values = Arc::new(StringArray::from(vec!["apple", "banana"]));
+ let left = DictionaryArray::<Int32Type>::new(left_keys, left_values);
+
+ // Right carries a logical null via a null VALUE referenced at row 1.
+ let right_keys = Int32Array::from(vec![Some(0_i32), Some(1), Some(0)]);
+ let right_values = Arc::new(StringArray::from(vec![Some("apple"),
None]));
+ let right = DictionaryArray::<Int32Type>::new(right_keys,
right_values);
+
+ let cmp = make_comparator(&left, &right,
SortOptions::default()).unwrap();
+ assert_eq!(cmp(0, 0), Ordering::Equal); // apple vs apple
+ assert_eq!(cmp(1, 0), Ordering::Less); // null-key < apple
(nulls_first)
+ assert_eq!(cmp(0, 1), Ordering::Greater); // apple > null-value
(nulls_first)
+ assert_eq!(cmp(1, 1), Ordering::Equal); // null-key vs null-value
+ }
Review Comment:
Could we add a few more cases here? This test only uses
`SortOptions::default()`, so nulls with `descending: true` or `nulls_first:
false` don't go through the new path in any test. `FixedSizeBinary` values with
nulls aren't covered either. `LargeUtf8`, `Binary` and `LargeBinary` each have
their own match arm in `compare_dict` but no test. A `#[should_panic(expected =
"index out of bounds")]` test like `test_union_out_of_bounds` would check the
out-of-bounds behavior from the comment on line 266.
##########
arrow-cmp/src/lib.rs:
##########
@@ -234,6 +253,93 @@ fn compare_dict<K: ArrowDictionaryKeyType>(
Ok(f)
}
+#[inline]
+fn compare_dict_bytes<K: ArrowDictionaryKeyType, T: ByteArrayType>(
+ left: &DictionaryArray<K>,
+ right: &DictionaryArray<K>,
+ opts: SortOptions,
+) -> DynComparator {
+ let left_values = left.values().as_bytes::<T>().clone();
+ let right_values = right.values().as_bytes::<T>().clone();
+ let left_keys = left.keys().values().clone();
+ let right_keys = right.keys().values().clone();
+ // SAFETY: row indices are in-bounds per DynComparator contract; keys are
+ // in-bounds for values per DictionaryArray::try_new.
+ if left.logical_null_count() == 0 && right.logical_null_count() == 0 {
+ return if opts.descending {
+ Box::new(move |left_row, right_row| unsafe {
+ let left_bytes: &[u8] = left_values
+
.value_unchecked(left_keys.get_unchecked(left_row).as_usize())
+ .as_ref();
+ let right_bytes: &[u8] = right_values
+
.value_unchecked(right_keys.get_unchecked(right_row).as_usize())
+ .as_ref();
+ right_bytes.cmp(left_bytes)
+ })
+ } else {
+ Box::new(move |left_row, right_row| unsafe {
+ let left_bytes: &[u8] = left_values
+
.value_unchecked(left_keys.get_unchecked(left_row).as_usize())
+ .as_ref();
+ let right_bytes: &[u8] = right_values
+
.value_unchecked(right_keys.get_unchecked(right_row).as_usize())
+ .as_ref();
+ left_bytes.cmp(right_bytes)
+ })
+ };
+ }
+ // SAFETY: same invariants as above.
+ compare(left, right, opts, move |left_row, right_row| unsafe {
+ let left_bytes: &[u8] = left_values
+ .value_unchecked(left_keys.get_unchecked(left_row).as_usize())
+ .as_ref();
+ let right_bytes: &[u8] = right_values
+ .value_unchecked(right_keys.get_unchecked(right_row).as_usize())
+ .as_ref();
+ left_bytes.cmp(right_bytes)
+ })
+}
Review Comment:
What do you think about dropping `#[inline]` and the hand-built no-null
closures, and calling `compare` once?
`compare` already returns the bare closure when neither side has logical
nulls ([`(None, None) =>
Box::new(cmp)`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L116-L117)).
`DESCENDING` is a const generic, so the `.reverse()` is resolved at compile
time
([here](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-cmp/src/lib.rs#L106-L109)).
I compared the code generated on aarch64 with rustc 1.99.0 in three builds:
the default release profile, fat LTO with `codegen-units = 1`, and thin LTO
with `codegen-units = 1`. In all three, the ascending closure in this PR and
the `(None, None)` closure from `compare_impl::<_, false, _>` are the same 33
instructions. The descending pair differs only in operand order and the final
condition code:
```
; closure#0 in this PR (operands swapped) ; compare_impl::<_, true, _>,
(None, None) arm
subs x19, x12, x9 subs x19, x9, x12
csel x2, x12, x9, lo csel x2, x9, x12, lo
...
csinv w0, w8, wzr, pl csinv w0, w8, wzr, ge
```
`#[inline]` doesn't change the output either. Both functions are private and
generic, each has a single call site, and they run once when the comparator is
built. The listings with and without the attribute are identical in all three
builds. The per-row closure is called through the `Box<dyn Fn>` vtable, so no
attribute on the constructor can inline it into the caller. The
[Reference](https://github.com/rust-lang/reference/blob/2dafab96d919d6ab2884488e970bf95610ff83eb/src/attributes/codegen.md?plain=1#L72-L73)
also says "In every form the attribute is a hint. The compiler may ignore it."
A single `compare` call also avoids a second pass over the keys. When the
dictionary values contain nulls,
[`logical_null_count`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-array/src/array/dictionary_array.rs#L769-L786)
scans every key, and then `compare` scans them again in
[`logical_nulls`](https://github.com/apache/arrow-rs/blob/72887ed91c80002bb7ec755ce1e964bab07bebea/arrow-array/src/array/dictionary_array.rs#L748-L767).
The suggestion also uses checked key lookups. The comment on line 266
explains why.
```suggestion
fn compare_dict_bytes<K: ArrowDictionaryKeyType, T: ByteArrayType>(
left: &DictionaryArray<K>,
right: &DictionaryArray<K>,
opts: SortOptions,
) -> DynComparator {
let left_values = left.values().as_bytes::<T>().clone();
let right_values = right.values().as_bytes::<T>().clone();
let left_keys = left.keys().values().clone();
let right_keys = right.keys().values().clone();
compare(left, right, opts, move |left_row, right_row| {
let left_key = left_keys[left_row].as_usize();
let right_key = right_keys[right_row].as_usize();
debug_assert!(left_key < left_values.len() && right_key <
right_values.len());
// SAFETY: `compare` calls this only for rows that are valid on both
sides, and
// `DictionaryArray` guarantees that every valid key indexes into
its values.
let (left_bytes, right_bytes): (&[u8], &[u8]) = unsafe {
(
left_values.value_unchecked(left_key).as_ref(),
right_values.value_unchecked(right_key).as_ref(),
)
};
left_bytes.cmp(right_bytes)
})
}
```
--
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]
