tustvold commented on code in PR #5792:
URL: https://github.com/apache/arrow-rs/pull/5792#discussion_r1611302298
##########
arrow-ord/src/ord.rs:
##########
@@ -20,104 +20,339 @@
use arrow_array::cast::AsArray;
use arrow_array::types::*;
use arrow_array::*;
-use arrow_buffer::ArrowNativeType;
-use arrow_schema::ArrowError;
+use arrow_buffer::{ArrowNativeType, NullBuffer};
+use arrow_schema::{ArrowError, SortOptions};
use std::cmp::Ordering;
/// Compare the values at two arbitrary indices in two arrays.
pub type DynComparator = Box<dyn Fn(usize, usize) -> Ordering + Send + Sync>;
-fn compare_primitive<T: ArrowPrimitiveType>(left: &dyn Array, right: &dyn
Array) -> DynComparator
+/// If parent sort order is descending we need to invert the value of
nulls_first so that
+/// when the parent is sorted based on the produced ranks, nulls are still
ordered correctly
+fn child_opts(opts: SortOptions) -> SortOptions {
+ SortOptions {
+ descending: false,
+ nulls_first: opts.nulls_first != opts.descending,
+ }
+}
+
+fn compare<A, F>(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator
where
- T::Native: ArrowNativeTypeOp,
+ A: Array + Clone,
+ F: Fn(usize, usize) -> Ordering + Send + Sync + 'static,
{
- let left = left.as_primitive::<T>().clone();
- let right = right.as_primitive::<T>().clone();
- Box::new(move |i, j| left.value(i).compare(right.value(j)))
+ let l = l.logical_nulls().filter(|x| x.null_count() > 0);
+ let r = r.logical_nulls().filter(|x| x.null_count() > 0);
+ match (opts.nulls_first, opts.descending) {
+ (true, true) => compare_impl::<true, true, _>(l, r, cmp),
+ (true, false) => compare_impl::<true, false, _>(l, r, cmp),
+ (false, true) => compare_impl::<false, true, _>(l, r, cmp),
+ (false, false) => compare_impl::<false, false, _>(l, r, cmp),
+ }
}
-fn compare_boolean(left: &dyn Array, right: &dyn Array) -> DynComparator {
- let left: BooleanArray = left.as_boolean().clone();
- let right: BooleanArray = right.as_boolean().clone();
+fn compare_impl<const NULLS_FIRST: bool, const DESCENDING: bool, F>(
Review Comment:
Avoids conditional branches in the comparator closure
--
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]
