diff --git a/arrow/src/compute/kernels/sort.rs b/arrow/src/compute/kernels/sort.rs index 0a02167dc528..ee3af71142d4 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::(values, v, n, &options, limit) } + DataType::Float32 => { + sort_list::(values, v, n, &options, limit) + } + DataType::Float64 => { + sort_list::(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::(values, v, n, &options, limit) } + DataType::Float32 => { + sort_list::(values, v, n, &options, limit) + } + DataType::Float64 => { + sort_list::(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::(values, v, n, &options, limit) } + DataType::Float32 => { + sort_list::(values, v, n, &options, limit) + } + DataType::Float64 => { + sort_list::(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::( + 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::( + 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::( vec![ Some(vec![Some(1), Some(0)]),