From 141b47487877893ee49834507d78b00ab0c7cc56 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Tue, 17 Mar 2026 11:33:30 -0400 Subject: [PATCH 1/3] Add array_sort benchmark --- datafusion/functions-nested/Cargo.toml | 4 + .../functions-nested/benches/array_sort.rs | 135 ++++++++++++++++++ datafusion/functions-nested/src/sort.rs | 2 +- 3 files changed, 140 insertions(+), 1 deletion(-) create mode 100644 datafusion/functions-nested/benches/array_sort.rs diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index 5fce3e854eb33..2ce9532a22ee7 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/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" diff --git a/datafusion/functions-nested/benches/array_sort.rs b/datafusion/functions-nested/benches/array_sort.rs new file mode 100644 index 0000000000000..edd95f038ceac --- /dev/null +++ b/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 = (0..total_values as i32).collect(); + values.shuffle(&mut rng); + + let values = Arc::new(Int32Array::from(values)); + let offsets: Vec = (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::>(), + )) + } 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 = (0..total_values).collect(); + indices.shuffle(&mut rng); + let string_values: Vec = + indices.iter().map(|i| format!("value_{:06}", i)).collect(); + let values = Arc::new(StringArray::from(string_values)); + + let offsets: Vec = (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(&[array.clone()]).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(&[array.clone()]).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(&[array.clone()]).unwrap()); + }); + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_array_sort); +criterion_main!(benches); diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs index cbe101f111b26..7b65dd96b3021 100644 --- a/datafusion/functions-nested/src/sort.rs +++ b/datafusion/functions-nested/src/sort.rs @@ -151,7 +151,7 @@ impl ScalarUDFImpl for ArraySort { } } -fn array_sort_inner(args: &[ArrayRef]) -> Result { +pub fn array_sort_inner(args: &[ArrayRef]) -> Result { if args.is_empty() || args.len() > 3 { return exec_err!("array_sort expects one to three arguments"); } From f51e313146a282c8e81bf4b0b7f97fbef12b4dcc Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Tue, 17 Mar 2026 14:28:04 -0400 Subject: [PATCH 2/3] Optimize array_sort for short string arrays --- datafusion/functions-nested/src/sort.rs | 166 ++++++++++++++---- datafusion/sqllogictest/test_files/array.slt | 16 ++ .../source/user-guide/sql/scalar_functions.md | 4 +- 3 files changed, 150 insertions(+), 36 deletions(-) diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs index 7b65dd96b3021..260786abab408 100644 --- a/datafusion/functions-nested/src/sort.rs +++ b/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)] @@ -208,55 +211,150 @@ fn array_sort_generic( list_array: &GenericListArray, field: FieldRef, sort_options: Option, +) -> Result { + 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( + list_array: &GenericListArray, + field: FieldRef, + sort_options: Option, ) -> Result { 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::::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( + list_array: &GenericListArray, + field: FieldRef, + sort_options: Option, +) -> Result { + 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 = 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 = 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::>(); + let sorted_values = if indices.is_empty() { + values.slice(0, 0) + } else { + take_by_indices::(values, &indices)? + }; + + Ok(Arc::new(GenericListArray::::try_new( + field, + OffsetBuffer::::new(new_offsets.into()), + sorted_values, + list_array.nulls().cloned(), + )?)) +} - let list_arr = if elements.is_empty() { - GenericListArray::::new_null(field, row_count) +/// Select elements from `values` at the given `indices` using `compute::take`. +fn take_by_indices( + values: &ArrayRef, + indices: &[OffsetSize], +) -> Result { + let indices_array: ArrayRef = if OffsetSize::IS_LARGE { + Arc::new(UInt64Array::from( + indices + .iter() + .map(|i| i.as_usize() as u64) + .collect::>(), + )) } else { - GenericListArray::::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::>(), + )) }; - Ok(Arc::new(list_arr)) + Ok(compute::take(values.as_ref(), &indices_array, None)?) } fn order_desc(modifier: &str) -> Result { diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index 83e9c9cc9c409..cf924aa4c6468 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/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)]) ---- diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md index 254151c2c20eb..0c0a1d40ea889 100644 --- a/docs/source/user-guide/sql/scalar_functions.md +++ b/docs/source/user-guide/sql/scalar_functions.md @@ -4182,8 +4182,8 @@ array_sort(array, desc, nulls_first) #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. -- **desc**: Whether to sort in descending order(`ASC` or `DESC`). -- **nulls_first**: Whether to sort nulls first(`NULLS FIRST` or `NULLS LAST`). +- **desc**: Whether to sort in ascending (`ASC`) or descending (`DESC`) order. The default is `ASC`. +- **nulls_first**: Whether to sort nulls first (`NULLS FIRST`) or last (`NULLS LAST`). The default is `NULLS FIRST`. #### Example From aa49a4d1bd7143fb6377beadf6c4a9493f484b9b Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Tue, 17 Mar 2026 14:43:52 -0400 Subject: [PATCH 3/3] Fix clippy --- datafusion/functions-nested/benches/array_sort.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/functions-nested/benches/array_sort.rs b/datafusion/functions-nested/benches/array_sort.rs index edd95f038ceac..bbf3ac6cd2bca 100644 --- a/datafusion/functions-nested/benches/array_sort.rs +++ b/datafusion/functions-nested/benches/array_sort.rs @@ -70,7 +70,7 @@ fn create_string_list_array(num_rows: usize, elements_per_row: usize) -> ArrayRe let mut indices: Vec = (0..total_values).collect(); indices.shuffle(&mut rng); let string_values: Vec = - indices.iter().map(|i| format!("value_{:06}", i)).collect(); + indices.iter().map(|i| format!("value_{i:06}")).collect(); let values = Arc::new(StringArray::from(string_values)); let offsets: Vec = (0..=num_rows) @@ -98,7 +98,7 @@ fn bench_array_sort(c: &mut Criterion) { &elements_per_row, |b, _| { b.iter(|| { - black_box(array_sort_inner(&[array.clone()]).unwrap()); + black_box(array_sort_inner(std::slice::from_ref(&array)).unwrap()); }); }, ); @@ -109,7 +109,7 @@ fn bench_array_sort(c: &mut Criterion) { let array = create_int32_list_array(NUM_ROWS, 50, true); group.bench_function("int32_with_nulls", |b| { b.iter(|| { - black_box(array_sort_inner(&[array.clone()]).unwrap()); + black_box(array_sort_inner(std::slice::from_ref(&array)).unwrap()); }); }); } @@ -122,7 +122,7 @@ fn bench_array_sort(c: &mut Criterion) { &elements_per_row, |b, _| { b.iter(|| { - black_box(array_sort_inner(&[array.clone()]).unwrap()); + black_box(array_sort_inner(std::slice::from_ref(&array)).unwrap()); }); }, );