From 1ef55992845396efef875f32096773c7a2f6bccc Mon Sep 17 00:00:00 2001
From: FANNG <sandflee.fang@gmail.com>
Date: Sun, 13 Nov 2022 21:06:26 +0800
Subject: [PATCH] Fixed descending ordering when specify nulls first (#1286)

---
 src/compute/merge_sort/mod.rs  | 13 +++++---
 tests/it/compute/merge_sort.rs | 57 ++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/src/compute/merge_sort/mod.rs b/src/compute/merge_sort/mod.rs
index 2776d58266e..598ba66d318 100644
--- a/src/compute/merge_sort/mod.rs
+++ b/src/compute/merge_sort/mod.rs
@@ -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
@@ -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;
diff --git a/tests/it/compute/merge_sort.rs b/tests/it/compute/merge_sort.rs
index 3adabb803a6..54b90233e53 100644
--- a/tests/it/compute/merge_sort.rs
+++ b/tests/it/compute/merge_sort.rs
@@ -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]);