From 3bc7f805f87d6684a3a10a09391135af1ce63abc Mon Sep 17 00:00:00 2001
From: Michael Edwards <medwards@walledcity.ca>
Date: Tue, 8 Jun 2021 22:55:17 +0200
Subject: [PATCH] Sort by float lists (#420)

---
 arrow/src/compute/kernels/sort.rs | 64 +++++++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/arrow/src/compute/kernels/sort.rs b/arrow/src/compute/kernels/sort.rs
index 7cd463d6ebb9..4723efc118a9 100644
--- a/arrow/src/compute/kernels/sort.rs
+++ b/arrow/src/compute/kernels/sort.rs
@@ -274,6 +274,12 @@ pub fn sort_to_indices(
             DataType::UInt64 => {
                 sort_list::<i32, UInt64Type>(values, v, n, &options, limit)
             }
+            DataType::Float32 => {
+                sort_list::<i32, Float32Type>(values, v, n, &options, limit)
+            }
+            DataType::Float64 => {
+                sort_list::<i32, Float64Type>(values, v, n, &options, limit)
+            }
             t => Err(ArrowError::ComputeError(format!(
                 "Sort not supported for list type {:?}",
                 t
@@ -294,6 +300,12 @@ pub fn sort_to_indices(
             DataType::UInt64 => {
                 sort_list::<i64, UInt64Type>(values, v, n, &options, limit)
             }
+            DataType::Float32 => {
+                sort_list::<i64, Float32Type>(values, v, n, &options, limit)
+            }
+            DataType::Float64 => {
+                sort_list::<i64, Float64Type>(values, v, n, &options, limit)
+            }
             t => Err(ArrowError::ComputeError(format!(
                 "Sort not supported for list type {:?}",
                 t
@@ -314,6 +326,12 @@ pub fn sort_to_indices(
             DataType::UInt64 => {
                 sort_list::<i32, UInt64Type>(values, v, n, &options, limit)
             }
+            DataType::Float32 => {
+                sort_list::<i32, Float32Type>(values, v, n, &options, limit)
+            }
+            DataType::Float64 => {
+                sort_list::<i32, Float64Type>(values, v, n, &options, limit)
+            }
             t => Err(ArrowError::ComputeError(format!(
                 "Sort not supported for list type {:?}",
                 t
@@ -2154,6 +2172,52 @@ mod tests {
             Some(1),
         );
 
+        test_sort_list_arrays::<Float32Type>(
+            vec![
+                Some(vec![Some(1.0), Some(0.0)]),
+                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
+                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
+                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
+                Some(vec![Some(1.0), Some(1.0)]),
+            ],
+            Some(SortOptions {
+                descending: false,
+                nulls_first: false,
+            }),
+            None,
+            vec![
+                Some(vec![Some(1.0), Some(0.0)]),
+                Some(vec![Some(1.0), Some(1.0)]),
+                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
+                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
+                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
+            ],
+            None,
+        );
+
+        test_sort_list_arrays::<Float64Type>(
+            vec![
+                Some(vec![Some(1.0), Some(0.0)]),
+                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
+                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
+                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
+                Some(vec![Some(1.0), Some(1.0)]),
+            ],
+            Some(SortOptions {
+                descending: false,
+                nulls_first: false,
+            }),
+            None,
+            vec![
+                Some(vec![Some(1.0), Some(0.0)]),
+                Some(vec![Some(1.0), Some(1.0)]),
+                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
+                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
+                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
+            ],
+            None,
+        );
+
         test_sort_list_arrays::<Int32Type>(
             vec![
                 Some(vec![Some(1), Some(0)]),