Fix array_sort for empty record batch (#290) (apache#15149) v48

Author: joroKr21

datafusion/functions-nested/src/sort.rs

@@ -70,6 +70,7 @@ impl ScalarUDFImpl for ArraySort {
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
         match &arg_types[0] {
+            DataType::Null => Ok(DataType::Null),
             List(field) | FixedSizeList(field, _) => Ok(List(Arc::new(
                 Field::new_list_field(field.data_type().clone(), true),
             ))),
@@ -141,6 +142,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 => {
@@ -161,12 +172,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(Arc::clone(&args[0]));
-    }
-
     let mut array_lengths = vec![];
     let mut arrays = vec![];
     let mut valid = BooleanBufferBuilder::new(row_count);

datafusion/sqllogictest/test_files/array.slt

@@ -2209,11 +2209,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