perf: Optimize array_sort for arrays of non-primitive types

datafusion/functions-nested/Cargo.toml

@@ -109,3 +109,7 @@ name = "array_to_string"
 [[bench]]
 harness = false
 name = "array_position"
+
+[[bench]]
+harness = false
+name = "array_sort"

datafusion/functions-nested/benches/array_sort.rs

@@ -0,0 +1,135 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, ListArray, StringArray};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::datatypes::{DataType, Field};
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_functions_nested::sort::array_sort_inner;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use rand::seq::SliceRandom;
+
+const SEED: u64 = 42;
+const NUM_ROWS: usize = 8192;
+
+fn create_int32_list_array(
+    num_rows: usize,
+    elements_per_row: usize,
+    with_nulls: bool,
+) -> ArrayRef {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let total_values = num_rows * elements_per_row;
+
+    let mut values: Vec<i32> = (0..total_values as i32).collect();
+    values.shuffle(&mut rng);
+
+    let values = Arc::new(Int32Array::from(values));
+    let offsets: Vec<i32> = (0..=num_rows)
+        .map(|i| (i * elements_per_row) as i32)
+        .collect();
+
+    let nulls = if with_nulls {
+        // Every 10th row is null
+        Some(NullBuffer::from(
+            (0..num_rows).map(|i| i % 10 != 0).collect::<Vec<bool>>(),
+        ))
+    } else {
+        None
+    };
+
+    Arc::new(ListArray::new(
+        Arc::new(Field::new("item", DataType::Int32, true)),
+        OffsetBuffer::new(offsets.into()),
+        values,
+        nulls,
+    ))
+}
+
+fn create_string_list_array(num_rows: usize, elements_per_row: usize) -> ArrayRef {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let total_values = num_rows * elements_per_row;
+
+    let mut indices: Vec<usize> = (0..total_values).collect();
+    indices.shuffle(&mut rng);
+    let string_values: Vec<String> =
+        indices.iter().map(|i| format!("value_{i:06}")).collect();
+    let values = Arc::new(StringArray::from(string_values));
+
+    let offsets: Vec<i32> = (0..=num_rows)
+        .map(|i| (i * elements_per_row) as i32)
+        .collect();
+
+    Arc::new(ListArray::new(
+        Arc::new(Field::new("item", DataType::Utf8, true)),
+        OffsetBuffer::new(offsets.into()),
+        values,
+        None,
+    ))
+}
+
+/// Vary elements_per_row over [5, 20, 100, 1000]: for small arrays, per-row
+/// overhead dominates, whereas for larger arrays the sort kernel dominates.
+fn bench_array_sort(c: &mut Criterion) {
+    let mut group = c.benchmark_group("array_sort");
+
+    // Int32 arrays
+    for &elements_per_row in &[5, 20, 100, 1000] {
+        let array = create_int32_list_array(NUM_ROWS, elements_per_row, false);
+        group.bench_with_input(
+            BenchmarkId::new("int32", elements_per_row),
+            &elements_per_row,
+            |b, _| {
+                b.iter(|| {
+                    black_box(array_sort_inner(std::slice::from_ref(&array)).unwrap());
+                });
+            },
+        );
+    }
+
+    // Int32 with nulls in the outer list (10% null rows), single size
+    {
+        let array = create_int32_list_array(NUM_ROWS, 50, true);
+        group.bench_function("int32_with_nulls", |b| {
+            b.iter(|| {
+                black_box(array_sort_inner(std::slice::from_ref(&array)).unwrap());
+            });
+        });
+    }
+
+    // String arrays
+    for &elements_per_row in &[5, 20, 100, 1000] {
+        let array = create_string_list_array(NUM_ROWS, elements_per_row);
+        group.bench_with_input(
+            BenchmarkId::new("string", elements_per_row),
+            &elements_per_row,
+            |b, _| {
+                b.iter(|| {
+                    black_box(array_sort_inner(std::slice::from_ref(&array)).unwrap());
+                });
+            },
+        );
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_array_sort);
+criterion_main!(benches);

datafusion/functions-nested/src/sort.rs

@@ -18,7 +18,10 @@
 //! [`ScalarUDFImpl`] definitions for array_sort function.
 
 use crate::utils::make_scalar_function;
-use arrow::array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait, new_null_array};
+use arrow::array::{
+    Array, ArrayRef, GenericListArray, OffsetSizeTrait, UInt32Array, UInt64Array,
+    new_null_array,
+};
 use arrow::buffer::OffsetBuffer;
 use arrow::compute::SortColumn;
 use arrow::datatypes::{DataType, FieldRef};
@@ -67,11 +70,11 @@ make_udf_expr_and_func!(
     ),
     argument(
         name = "desc",
-        description = "Whether to sort in descending order(`ASC` or `DESC`)."
+        description = "Whether to sort in ascending (`ASC`) or descending (`DESC`) order. The default is `ASC`."
     ),
     argument(
         name = "nulls_first",
-        description = "Whether to sort nulls first(`NULLS FIRST` or `NULLS LAST`)."
+        description = "Whether to sort nulls first (`NULLS FIRST`) or last (`NULLS LAST`). The default is `NULLS FIRST`."
     )
 )]
 #[derive(Debug, PartialEq, Eq, Hash)]
@@ -151,7 +154,7 @@ impl ScalarUDFImpl for ArraySort {
     }
 }
 
-fn array_sort_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+pub fn array_sort_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
     if args.is_empty() || args.len() > 3 {
         return exec_err!("array_sort expects one to three arguments");
     }
@@ -208,55 +211,150 @@ fn array_sort_generic<OffsetSize: OffsetSizeTrait>(
     list_array: &GenericListArray<OffsetSize>,
     field: FieldRef,
     sort_options: Option<SortOptions>,
+) -> Result<ArrayRef> {
+    let values = list_array.values();
+
+    if values.data_type().is_primitive() {
+        array_sort_direct(list_array, field, sort_options)
+    } else {
+        array_sort_batch_indices(list_array, field, sort_options)
+    }
+}
+
+/// Sort each row using `compute::sort()` and concatenate the results.
+///
+/// This is efficient for primitive element types because Arrow's sort kernel
+/// does the sorting in-place.
+fn array_sort_direct<OffsetSize: OffsetSizeTrait>(
+    list_array: &GenericListArray<OffsetSize>,
+    field: FieldRef,
+    sort_options: Option<SortOptions>,
 ) -> Result<ArrayRef> {
     let row_count = list_array.len();
+    let values = list_array.values();
 
-    let mut array_lengths = vec![];
-    let mut arrays = vec![];
+    let mut array_lengths = Vec::with_capacity(row_count);
+    let mut sorted_arrays = Vec::with_capacity(row_count);
     for i in 0..row_count {
         if list_array.is_null(i) {
             array_lengths.push(0);
         } else {
             let arr_ref = list_array.value(i);
+            let sorted = compute::sort(arr_ref.as_ref(), sort_options)?;
+            array_lengths.push(sorted.len());
+            sorted_arrays.push(sorted);
+        }
+    }
 
-            // arrow sort kernel does not support Structs, so use
+    let sorted_values: ArrayRef = if sorted_arrays.is_empty() {
+        values.slice(0, 0)
+    } else {
+        let elements: Vec<&dyn Array> =
+            sorted_arrays.iter().map(|a| a.as_ref()).collect();
+        Arc::new(compute::concat(&elements)?)
+    };
+
+    Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
+        field,
+        OffsetBuffer::from_lengths(array_lengths),
+        sorted_values,
+        list_array.nulls().cloned(),
+    )?))
+}
+
+/// Sort each row by collecting sort indices, then materialize with a single
+/// `take()` at the end.
+///
+/// This is efficient for non-primitive element types because Arrow's sort
+/// kernel would internally call `sort_to_indices()` + `take()` per row anyway.
+/// Batching into a single `take()` avoids N per-row allocations and the final
+/// `concat()`.
+fn array_sort_batch_indices<OffsetSize: OffsetSizeTrait>(
+    list_array: &GenericListArray<OffsetSize>,
+    field: FieldRef,
+    sort_options: Option<SortOptions>,
+) -> Result<ArrayRef> {
+    let row_count = list_array.len();
+    let values = list_array.values();
+    let offsets = list_array.offsets();
+
+    let total_values = offsets[row_count].as_usize() - offsets[0].as_usize();
+    let mut indices: Vec<OffsetSize> = Vec::with_capacity(total_values);
+    let mut new_offsets = Vec::with_capacity(row_count + 1);
+    new_offsets.push(OffsetSize::usize_as(0));
+
+    let is_struct = matches!(values.data_type(), DataType::Struct(_));
+
+    for (row_index, window) in offsets.windows(2).enumerate() {
+        let start = window[0];
+        let end = window[1];
+
+        if list_array.is_null(row_index) {
+            new_offsets.push(new_offsets[row_index]);
+            continue;
+        }
+
+        let len = (end - start).as_usize();
+        if len <= 1 {
+            // 0 or 1 elements: already sorted, push identity indices
+            indices.extend((start.as_usize()..end.as_usize()).map(OffsetSize::usize_as));
+        } else {
+            let sliced = values.slice(start.as_usize(), len);
+            // Arrow's sort kernel does not support Struct arrays, so use
             // lexsort_to_indices instead:
             // https://github.com/apache/arrow-rs/issues/6911#issuecomment-2562928843
-            let sorted_array = match arr_ref.data_type() {
-                DataType::Struct(_) => {
-                    let sort_columns: Vec<SortColumn> = vec![SortColumn {
-                        values: Arc::clone(&arr_ref),
-                        options: sort_options,
-                    }];
-                    let indices = compute::lexsort_to_indices(&sort_columns, None)?;
-                    compute::take(arr_ref.as_ref(), &indices, None)?
-                }
-                _ => {
-                    let arr_ref = arr_ref.as_ref();
-                    compute::sort(arr_ref, sort_options)?
-                }
+            let sorted_indices = if is_struct {
+                let sort_columns = vec![SortColumn {
+                    values: sliced,
+                    options: sort_options,
+                }];
+                compute::lexsort_to_indices(&sort_columns, None)?
+            } else {
+                compute::sort_to_indices(&sliced, sort_options, None)?
             };
-            array_lengths.push(sorted_array.len());
-            arrays.push(sorted_array);
+            for &local_idx in sorted_indices.values() {
+                indices.push(start + OffsetSize::usize_as(local_idx as usize));
+            }
         }
+
+        new_offsets.push(new_offsets[row_index] + (end - start));
     }
 
-    let elements = arrays
-        .iter()
-        .map(|a| a.as_ref())
-        .collect::<Vec<&dyn Array>>();
+    let sorted_values = if indices.is_empty() {
+        values.slice(0, 0)
+    } else {
+        take_by_indices::<OffsetSize>(values, &indices)?
+    };
+
+    Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
+        field,
+        OffsetBuffer::<OffsetSize>::new(new_offsets.into()),
+        sorted_values,
+        list_array.nulls().cloned(),
+    )?))
+}
 
-    let list_arr = if elements.is_empty() {
-        GenericListArray::<OffsetSize>::new_null(field, row_count)
+/// Select elements from `values` at the given `indices` using `compute::take`.
+fn take_by_indices<OffsetSize: OffsetSizeTrait>(
+    values: &ArrayRef,
+    indices: &[OffsetSize],
+) -> Result<ArrayRef> {
+    let indices_array: ArrayRef = if OffsetSize::IS_LARGE {
+        Arc::new(UInt64Array::from(
+            indices
+                .iter()
+                .map(|i| i.as_usize() as u64)
+                .collect::<Vec<_>>(),
+        ))
     } else {
-        GenericListArray::<OffsetSize>::new(
-            field,
-            OffsetBuffer::from_lengths(array_lengths),
-            Arc::new(compute::concat(elements.as_slice())?),
-            list_array.nulls().cloned(),
-        )
+        Arc::new(UInt32Array::from(
+            indices
+                .iter()
+                .map(|i| i.as_usize() as u32)
+                .collect::<Vec<_>>(),
+        ))
     };
-    Ok(Arc::new(list_arr))
+    Ok(compute::take(values.as_ref(), &indices_array, None)?)
 }
 
 fn order_desc(modifier: &str) -> Result<bool> {

datafusion/sqllogictest/test_files/array.slt

@@ -2535,6 +2535,14 @@ select array_sort([]);
 ----
 []
 
+# empty-but-non-null string arrays should remain non-null, not become null
+query ?B
+select array_sort(column1), array_sort(column1) is null
+from (values (arrow_cast(make_array('b', 'a'), 'List(Utf8)')), (arrow_cast([], 'List(Utf8)'))) as t(column1);
+----
+[a, b] false
+[] false
+
 # test with null arguments
 query ?
 select array_sort(NULL);
@@ -2602,6 +2610,14 @@ from values (array_sort(arrow_cast([1, 3, 5, -5], 'FixedSizeList(4 x non-null In
 ----
 [-5, 1, 3, 5] List(non-null Int32)
 
+# arrays of strings
+query ???
+select array_sort(make_array('banana', 'apple', null, 'cherry')),
+       array_sort(make_array('banana', 'apple', null, 'cherry'), 'DESC', 'NULLS LAST'),
+       array_sort(make_array('banana', 'apple', null, 'cherry'), 'ASC', 'NULLS LAST');
+----
+[NULL, apple, banana, cherry] [cherry, banana, apple, NULL] [apple, banana, cherry, NULL]
+
 query ?
 select array_sort([struct('foo', 3), struct('foo', 1), struct('bar', 1)])
 ----