Skip to content

Commit

Permalink
Fixed descending ordering when specify nulls first (jorgecarleitao#1286)
Browse files Browse the repository at this point in the history
  • Loading branch information
FANNG1 authored and ritchie46 committed Mar 29, 2023
1 parent ec50cc8 commit 1ef5599
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 5 deletions.
13 changes: 8 additions & 5 deletions src/compute/merge_sort/mod.rs
Expand Up @@ -509,8 +509,14 @@ pub fn build_comparator_impl<'a>(
let descending = pairs[c].1.descending;
let null_first = pairs[c].1.nulls_first;
let (l_is_valid, r_is_valid, value_comparator) = &data[c];
let mut result = match ((l_is_valid)(left_row), (r_is_valid)(right_row)) {
(true, true) => (value_comparator)(left_row, right_row),
let result = match ((l_is_valid)(left_row), (r_is_valid)(right_row)) {
(true, true) => {
let result = (value_comparator)(left_row, right_row);
match descending {
true => result.reverse(),
false => result,
}
}
(false, true) => {
if null_first {
Ordering::Less
Expand All @@ -527,9 +533,6 @@ pub fn build_comparator_impl<'a>(
}
(false, false) => Ordering::Equal,
};
if descending {
result = result.reverse();
};
if result != Ordering::Equal {
// we found a relevant comparison => short-circuit and return it
return result;
Expand Down
57 changes: 57 additions & 0 deletions tests/it/compute/merge_sort.rs
Expand Up @@ -34,6 +34,63 @@ fn merge_u32() -> Result<()> {
Ok(())
}

#[test]
fn merge_null_first() -> Result<()> {
let a0: &dyn Array = &Int32Array::from(&[None, Some(0)]);
let a1: &dyn Array = &Int32Array::from(&[Some(2), Some(3)]);
let options = SortOptions {
descending: false,
nulls_first: true,
};
let arrays = vec![a0, a1];
let pairs = vec![(arrays.as_ref(), &options)];
let comparator = build_comparator(&pairs)?;
let result =
merge_sort_slices(once(&(0, 0, 2)), once(&(1, 0, 2)), &comparator).collect::<Vec<_>>();
assert_eq!(result, vec![(0, 0, 2), (1, 0, 2)]);

let a0: &dyn Array = &Int32Array::from(&[Some(0), None]);
let a1: &dyn Array = &Int32Array::from(&[Some(2), Some(3)]);
let options = SortOptions {
descending: false,
nulls_first: false,
};
let arrays = vec![a0, a1];
let pairs = vec![(arrays.as_ref(), &options)];
let comparator = build_comparator(&pairs)?;
let result =
merge_sort_slices(once(&(0, 0, 2)), once(&(1, 0, 2)), &comparator).collect::<Vec<_>>();
assert_eq!(result, vec![(0, 0, 1), (1, 0, 2), (0, 1, 1)]);

let a0: &dyn Array = &Int32Array::from(&[Some(0), None]);
let a1: &dyn Array = &Int32Array::from(&[Some(3), Some(2)]);
let options = SortOptions {
descending: true,
nulls_first: false,
};
let arrays = vec![a0, a1];
let pairs = vec![(arrays.as_ref(), &options)];
let comparator = build_comparator(&pairs)?;
let result =
merge_sort_slices(once(&(0, 0, 2)), once(&(1, 0, 2)), &comparator).collect::<Vec<_>>();
assert_eq!(result, vec![(1, 0, 2), (0, 0, 2)]);

let a0: &dyn Array = &Int32Array::from(&[None, Some(0)]);
let a1: &dyn Array = &Int32Array::from(&[Some(3), Some(2)]);
let options = SortOptions {
descending: true,
nulls_first: true,
};
let arrays = vec![a0, a1];
let pairs = vec![(arrays.as_ref(), &options)];
let comparator = build_comparator(&pairs)?;
let result =
merge_sort_slices(once(&(0, 0, 2)), once(&(1, 0, 2)), &comparator).collect::<Vec<_>>();
assert_eq!(result, vec![(0, 0, 1), (1, 0, 2), (0, 1, 1)]);

Ok(())
}

#[test]
fn merge_with_limit() -> Result<()> {
let a0: &dyn Array = &Int32Array::from_slice([0, 2, 4, 6, 8]);
Expand Down

0 comments on commit 1ef5599

Please sign in to comment.