Fix array_sort for empty record batch (#290)

Author: joroKr21

datafusion/functions-array/src/sort.rs

@@ -21,7 +21,6 @@ use crate::utils::make_scalar_function;
 use arrow::compute;
 use arrow_array::{Array, ArrayRef, ListArray};
 use arrow_buffer::{BooleanBufferBuilder, NullBuffer, OffsetBuffer};
-use arrow_schema::DataType::{FixedSizeList, LargeList, List};
 use arrow_schema::{DataType, Field, SortOptions};
 use datafusion_common::cast::{as_list_array, as_string_array};
 use datafusion_common::{exec_err, Result};
@@ -67,19 +66,9 @@ impl ScalarUDFImpl for ArraySort {
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         match &arg_types[0] {
-            List(field) | FixedSizeList(field, _) => Ok(List(Arc::new(Field::new(
-                "item",
-                field.data_type().clone(),
-                true,
-            )))),
-            LargeList(field) => Ok(LargeList(Arc::new(Field::new(
-                "item",
-                field.data_type().clone(),
-                true,
-            )))),
-            _ => exec_err!(
-                "Not reachable, data_type should be List, LargeList or FixedSizeList"
-            ),
+            DataType::Null => Ok(DataType::Null),
+            arg_type @ DataType::List(_) => Ok(arg_type.clone()),
+            arg_type => exec_err!("{} does not support type {arg_type}", self.name()),
         }
     }
 
@@ -98,6 +87,16 @@ pub fn array_sort_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
         return exec_err!("array_sort expects one to three arguments");
     }
 
+    if args[0].data_type().is_null() {
+        return Ok(Arc::clone(&args[0]));
+    }
+
+    let list_array = as_list_array(&args[0])?;
+    let row_count = list_array.len();
+    if row_count == 0 || list_array.value_type().is_null() {
+        return Ok(Arc::clone(&args[0]));
+    }
+
     let sort_option = match args.len() {
         1 => None,
         2 => {
@@ -118,12 +117,6 @@ pub fn array_sort_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
         _ => return exec_err!("array_sort expects 1 to 3 arguments"),
     };
 
-    let list_array = as_list_array(&args[0])?;
-    let row_count = list_array.len();
-    if row_count == 0 {
-        return Ok(args[0].clone());
-    }
-
     let mut array_lengths = vec![];
     let mut arrays = vec![];
     let mut valid = BooleanBufferBuilder::new(row_count);

datafusion/sqllogictest/test_files/array.slt

@@ -2038,11 +2038,16 @@ NULL
 [, 51, 52, 54, 55, 56, 57, 58, 59, 60]
 [61, 62, 63, 64, 65, 66, 67, 68, 69, 70]
 
-# test with empty array
+# test with empty table
 query ?
-select array_sort([]);
+select array_sort(column1, 'DESC', 'NULLS FIRST') from arrays_values where false;
 ----
-[]
+
+# test with empty array
+query ??
+select array_sort([]), array_sort(NULL);
+----
+[] NULL
 
 # test with empty row, the row that does not match the condition has row count 0
 statement ok