You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/22 01:46:18 UTC
(datafusion) branch main updated: perf: Optimize `array_sort` to count validity bits while compacting (#26453)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 83d2e9d12f perf: Optimize `array_sort` to count validity bits while compacting (#26453)
83d2e9d12f is described below
commit 83d2e9d12f66c61a1259c80f6df0c242299ad4b3
Author: Neil Conway <ne...@gmail.com>
AuthorDate: Thu Oct 22 00:54:58 2026 +0000
perf: Optimize `array_sort` to count validity bits while compacting (#26453)
## Which issue does this PR close?
- N/A
## Rationale for this change
This PR optimizes `array_sort` on arrays of primitive values, where at
least one element in the batch is NULL.
Previously, `array_sort` started by counting the NULL elements in the
row and then compacting the non-NULL elements to place NULLs either
first or last (depending on the specified sort order). Doing the
compaction already requires looking at every element's validity bit, so
we can make this more efficient by doing the compaction first and
deriving the null count as we go. When `NULLS FIRST`, we then need to
move the non-NULL values after the NULLs. Empirically, the cost of the
move is outweighed by the savings from avoiding a separate null count.
The compaction loop avoids branching on each element's validity bit:
when NULLs are scattered through a row, the CPU frequently mispredicts
that branch. Instead, write every value to the next output slot and
advance the write position only on valid values (which we can do without
a branch!). Per suggestion by @mbutrovich.
One consequence of the previous approach is that values in NULL slots
were always the zeroed. We don't attempt to preserve that behavior,
because Arrow allows NULL slots to contain arbitrary values.
Two alternatives were considered:
1. Using `SetBitCounter` from #25754; this was faster on some inputs and
slower on others.
2. Compacting values at the end of the buffer for `NULLS FIRST`, walking
elements backwards but avoiding the post-compaction move; this was
faster for small rows and slower for large rows.
Along the way, improve the `array_sort` benchmarks:
1. Add coverage for `NULLS LAST`
2. Cycle through several input arrays and make the distribution of NULLs
and array lengths more unpredictable, to avoid the branch predictor
being able to memorize the benchmark input.
## Benchmarks
x86, AMD EPYC-Milan:
int32_null_elements/5 -24.3%
int32_null_elements/20 -19.5%
int32_null_elements/100 -13.1%
int32_null_elements/1000 -8.9%
int32_null_elements_desc/5 -24.3%
int32_null_elements_desc/20 -19.1%
int32_null_elements_desc/100 -13.6%
int32_null_elements_desc/1000 -9.1%
int32_null_elements_nulls_last/5 -28.4%
int32_null_elements_nulls_last/20 -21.1%
int32_null_elements_nulls_last/100 -14.6%
int32_null_elements_nulls_last/1000 -9.4%
## What changes are included in this PR?
See above.
## What is the testing strategy for this PR?
Existing tests pass, new tests added.
## Are there any user-facing changes?
No.
---
datafusion/functions-nested/benches/array_sort.rs | 250 +++++++++-------------
datafusion/functions-nested/src/sort.rs | 96 +++++++--
2 files changed, 183 insertions(+), 163 deletions(-)
diff --git a/datafusion/functions-nested/benches/array_sort.rs b/datafusion/functions-nested/benches/array_sort.rs
index 2ceee91f4c..5cba6c6b93 100644
--- a/datafusion/functions-nested/benches/array_sort.rs
+++ b/datafusion/functions-nested/benches/array_sort.rs
@@ -18,123 +18,102 @@
use std::hint::black_box;
use std::sync::Arc;
-use arrow::array::{ArrayRef, BooleanBufferBuilder, Int32Array, ListArray, StringArray};
+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 criterion::{Bencher, BenchmarkId, Criterion, criterion_group, criterion_main};
use datafusion_common::ScalarValue;
use datafusion_common::config::ConfigOptions;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
use datafusion_functions_nested::sort::ArraySort;
-use rand::SeedableRng;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
+use rand::{Rng, SeedableRng};
const SEED: u64 = 42;
const NUM_ROWS: usize = 8192;
+/// Each case sorts the next of its arrays in each iteration, cycling through
+/// arrays that hold at least this many elements in total.
+const CYCLE_ELEMENTS: usize = 1 << 20;
-fn create_int32_list_array(
- num_rows: usize,
+/// The arrays a case cycles through, each built by `make` from its own seed:
+/// enough to hold `CYCLE_ELEMENTS` elements, and at least two.
+fn inputs(
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
- };
+ make: impl Fn(&mut StdRng) -> ArrayRef,
+) -> Vec<ArrayRef> {
+ let count = CYCLE_ELEMENTS.div_ceil(NUM_ROWS * elements_per_row).max(2);
+ (0..count as u64)
+ .map(|i| make(&mut StdRng::seed_from_u64(SEED + i)))
+ .collect()
+}
- Arc::new(ListArray::new(
- Arc::new(Field::new("item", DataType::Int32, true)),
- OffsetBuffer::new(offsets.into()),
- values,
- nulls,
- ))
+/// Offsets for `NUM_ROWS` rows of 1 to `2 * elements_per_row - 1` elements,
+/// `elements_per_row` on average.
+fn random_offsets(rng: &mut StdRng, elements_per_row: usize) -> OffsetBuffer<i32> {
+ OffsetBuffer::from_lengths(
+ (0..NUM_ROWS).map(|_| rng.random_range(1..2 * elements_per_row)),
+ )
+}
+
+/// Validity for `len` slots, each NULL with probability `null_rate`, or `None`
+/// if `null_rate` is zero.
+fn random_nulls(rng: &mut StdRng, len: usize, null_rate: f64) -> Option<NullBuffer> {
+ (null_rate > 0.0)
+ .then(|| NullBuffer::from_iter((0..len).map(|_| !rng.random_bool(null_rate))))
}
-/// Creates a ListArray where ~10% of elements within each row are null.
-fn create_int32_list_array_with_null_elements(
- num_rows: usize,
+/// A list of distinct `Int32` values in random order, where each row is NULL
+/// with probability `null_row_rate` and each element with probability
+/// `null_element_rate`.
+fn create_int32_list_array(
+ rng: &mut StdRng,
elements_per_row: usize,
+ null_row_rate: f64,
+ null_element_rate: f64,
) -> 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);
-
- // ~10% of elements are null
- let mut validity = BooleanBufferBuilder::new(total_values);
- for i in 0..total_values {
- validity.append(i % 10 != 0);
- }
- let null_buffer = NullBuffer::from(validity.finish());
-
- let values = Arc::new(Int32Array::new(values.into(), Some(null_buffer)));
- let offsets: Vec<i32> = (0..=num_rows)
- .map(|i| (i * elements_per_row) as i32)
- .collect();
+ let offsets = random_offsets(rng, elements_per_row);
+ let mut values: Vec<i32> = (0..offsets.last()).collect();
+ values.shuffle(rng);
+ let element_nulls = random_nulls(rng, values.len(), null_element_rate);
+ let list_nulls = random_nulls(rng, NUM_ROWS, null_row_rate);
Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Int32, true)),
- OffsetBuffer::new(offsets.into()),
- values,
- None,
+ offsets,
+ Arc::new(Int32Array::new(values.into(), element_nulls)),
+ list_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();
+fn create_string_list_array(rng: &mut StdRng, elements_per_row: usize) -> ArrayRef {
+ let offsets = random_offsets(rng, elements_per_row);
+ let mut indices: Vec<i32> = (0..offsets.last()).collect();
+ indices.shuffle(rng);
+ let values =
+ StringArray::from_iter_values(indices.iter().map(|i| format!("value_{i:06}")));
Arc::new(ListArray::new(
Arc::new(Field::new("item", DataType::Utf8, true)),
- OffsetBuffer::new(offsets.into()),
- values,
+ offsets,
+ Arc::new(values),
None,
))
}
-fn invoke_array_sort(udf: &ArraySort, array: &ArrayRef) -> ColumnarValue {
- invoke_array_sort_with_order(udf, array, None)
-}
-
-/// `order` is the optional second argument (`"ASC"` / `"DESC"`); `None` leaves
-/// `array_sort` with its default (ascending) ordering.
-fn invoke_array_sort_with_order(
+/// Calls `array_sort` on `array`, followed by `options`: the sort order
+/// (`"ASC"` / `"DESC"`), then where NULLs go (`"NULLS FIRST"` / `"NULLS LAST"`).
+/// Options left out take their defaults, ascending with NULLs first.
+fn invoke_array_sort(
udf: &ArraySort,
array: &ArrayRef,
- order: Option<&str>,
+ options: &[&str],
) -> ColumnarValue {
let mut args = vec![ColumnarValue::Array(Arc::clone(array))];
let mut arg_fields = vec![Field::new("arr", array.data_type().clone(), true).into()];
- if let Some(order) = order {
- args.push(ColumnarValue::Scalar(ScalarValue::from(order)));
- arg_fields.push(Field::new("order", DataType::Utf8, true).into());
+ for (name, &option) in ["order", "nulls"].into_iter().zip(options) {
+ args.push(ColumnarValue::Scalar(ScalarValue::from(option)));
+ arg_fields.push(Field::new(name, DataType::Utf8, true).into());
}
udf.invoke_with_args(ScalarFunctionArgs {
args,
@@ -146,91 +125,72 @@ fn invoke_array_sort_with_order(
.unwrap()
}
-/// Vary elements_per_row over [5, 20, 100, 1000]: for small arrays, per-row
-/// overhead dominates, whereas for larger arrays the sort kernel dominates.
+/// Sorts the next of `arrays` with `options` in each iteration.
+fn sort_each<'a>(
+ udf: &'a ArraySort,
+ arrays: &'a [ArrayRef],
+ options: &'a [&'a str],
+) -> impl FnMut(&mut Bencher) + 'a {
+ let mut arrays = arrays.iter().cycle();
+ move |b| b.iter(|| black_box(invoke_array_sort(udf, arrays.next().unwrap(), options)))
+}
+
+/// Vary the average 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");
let udf = ArraySort::new();
- // Int32 arrays
+ // Int32 arrays, ascending and descending: `sort_row` sorts ascending and
+ // reverses, so the descending case measures the cost of that extra pass.
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(
+ let arrays = inputs(elements_per_row, |rng| {
+ create_int32_list_array(rng, elements_per_row, 0.0, 0.0)
+ });
+ group.bench_function(
BenchmarkId::new("int32", elements_per_row),
- &elements_per_row,
- |b, _| {
- b.iter(|| {
- black_box(invoke_array_sort(&udf, &array));
- });
- },
+ sort_each(&udf, &arrays, &[]),
);
- }
-
- // Same shapes, descending: `sort_row` sorts ascending and reverses, so this
- // measures the cost of that extra pass.
- 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(
+ group.bench_function(
BenchmarkId::new("int32_desc", elements_per_row),
- &elements_per_row,
- |b, _| {
- b.iter(|| {
- black_box(invoke_array_sort_with_order(&udf, &array, Some("DESC")));
- });
- },
+ sort_each(&udf, &arrays, &["DESC"]),
);
}
// 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(invoke_array_sort(&udf, &array));
- });
- });
- }
-
- // Int32 with null elements (~10% of elements within rows are null)
- for &elements_per_row in &[5, 20, 100, 1000] {
- let array =
- create_int32_list_array_with_null_elements(NUM_ROWS, elements_per_row);
- group.bench_with_input(
- BenchmarkId::new("int32_null_elements", elements_per_row),
- &elements_per_row,
- |b, _| {
- b.iter(|| {
- black_box(invoke_array_sort(&udf, &array));
- });
- },
- );
+ let arrays = inputs(50, |rng| create_int32_list_array(rng, 50, 0.1, 0.0));
+ group.bench_function("int32_with_nulls", sort_each(&udf, &arrays, &[]));
}
+ // Int32 with null elements (10% of elements), placed first in ascending
+ // and descending order, and last in ascending order
for &elements_per_row in &[5, 20, 100, 1000] {
- let array =
- create_int32_list_array_with_null_elements(NUM_ROWS, elements_per_row);
- group.bench_with_input(
- BenchmarkId::new("int32_null_elements_desc", elements_per_row),
- &elements_per_row,
- |b, _| {
- b.iter(|| {
- black_box(invoke_array_sort_with_order(&udf, &array, Some("DESC")));
- });
- },
- );
+ let arrays = inputs(elements_per_row, |rng| {
+ create_int32_list_array(rng, elements_per_row, 0.0, 0.1)
+ });
+ let cases: [(&str, &[&str]); 3] = [
+ ("int32_null_elements", &[]),
+ ("int32_null_elements_desc", &["DESC"]),
+ ("int32_null_elements_nulls_last", &["ASC", "NULLS LAST"]),
+ ];
+ for (name, options) in cases {
+ group.bench_function(
+ BenchmarkId::new(name, elements_per_row),
+ sort_each(&udf, &arrays, options),
+ );
+ }
}
// 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(
+ let arrays = inputs(elements_per_row, |rng| {
+ create_string_list_array(rng, elements_per_row)
+ });
+ group.bench_function(
BenchmarkId::new("string", elements_per_row),
- &elements_per_row,
- |b, _| {
- b.iter(|| {
- black_box(invoke_array_sort(&udf, &array));
- });
- },
+ sort_each(&udf, &arrays, &[]),
);
}
diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs
index 585f3445ce..533f39f3c8 100644
--- a/datafusion/functions-nested/src/sort.rs
+++ b/datafusion/functions-nested/src/sort.rs
@@ -26,6 +26,7 @@ use arrow::array::{
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::datatypes::{ArrowNativeTypeOp, DataType};
use arrow::row::{RowConverter, Rows, SortField};
+use arrow::util::bit_iterator::BitIterator;
use arrow::{compute, compute::SortOptions, downcast_primitive_array};
use datafusion_common::cast::{as_large_list_array, as_list_array, as_string_array};
use datafusion_common::utils::{ListCoercion, offset_span};
@@ -335,28 +336,29 @@ fn sort_rows_with_nulls<N: ArrowNativeTypeOp, O: OffsetSizeTrait>(
continue;
}
- let null_count = src_nulls.slice(start, row_len).null_count();
- let valid_count = row_len - null_count;
-
- // Compact valid values directly into the target region of the output
- // buffer: after nulls (if nulls_first) or at the start (if nulls_last).
- let valid_offset = if nulls_first { null_count } else { 0 };
- let mut write_pos = out_start + valid_offset;
- for (i, &value) in src_values[start..end].iter().enumerate() {
- if src_nulls.is_valid(start + i) {
- out_values[write_pos] = value;
- write_pos += 1;
- }
+ // Compact valid values to the start of the row's output region and
+ // count them. The loop writes every value to the next slot but
+ // advances only past valid values, so it doesn't branch on each
+ // validity bit; that branch would be hard to predict when NULLs are
+ // scattered through the row.
+ let validity_bits =
+ BitIterator::new(src_nulls.validity(), src_nulls.offset() + start, row_len);
+ let mut write_pos = out_start;
+ for (&value, is_valid) in src_values[start..end].iter().zip(validity_bits) {
+ out_values[write_pos] = value;
+ write_pos += is_valid as usize;
}
+ let valid_count = write_pos - out_start;
+ let null_count = row_len - valid_count;
- sort_row(
- &mut out_values
- [out_start + valid_offset..out_start + valid_offset + valid_count],
- descending,
- );
+ sort_row(&mut out_values[out_start..write_pos], descending);
- // Build validity bits
+ // Build validity bits. With nulls_first, move the sorted values after
+ // the NULLs.
if nulls_first {
+ if null_count > 0 {
+ out_values.copy_within(out_start..write_pos, out_start + null_count);
+ }
validity.append_n(null_count, false);
validity.append_n(valid_count, true);
} else {
@@ -487,3 +489,61 @@ fn order_nulls_first(modifier: &str) -> Result<bool> {
),
}
}
+
+#[cfg(test)]
+mod tests {
+ use arrow::array::{Array, ArrayRef, Int32Array, ListArray, StringArray};
+ use arrow::buffer::OffsetBuffer;
+ use arrow::datatypes::{DataType, Field, Int32Type};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_sliced_list() -> datafusion_common::Result<()> {
+ crate::utils::tests::check_sliced_list_behavior(|input| {
+ super::array_sort_inner(&[Arc::clone(input)])
+ })
+ }
+
+ #[test]
+ fn test_sliced_values() -> datafusion_common::Result<()> {
+ // Slicing the child array makes its validity bitmap start partway
+ // through a byte.
+ let values = Int32Array::from(vec![
+ Some(0),
+ Some(0),
+ Some(0),
+ Some(3),
+ None,
+ Some(1),
+ Some(2),
+ None,
+ ]);
+ let input: ArrayRef = Arc::new(ListArray::new(
+ Arc::new(Field::new_list_field(DataType::Int32, true)),
+ OffsetBuffer::from_lengths([3, 2]),
+ Arc::new(values.slice(3, 5)),
+ None,
+ ));
+ for (nulls, expected) in [
+ (
+ "NULLS FIRST",
+ [vec![None, Some(1), Some(3)], vec![None, Some(2)]],
+ ),
+ (
+ "NULLS LAST",
+ [vec![Some(1), Some(3), None], vec![Some(2), None]],
+ ),
+ ] {
+ let result = super::array_sort_inner(&[
+ Arc::clone(&input),
+ Arc::new(StringArray::from(vec!["ASC"; 2])),
+ Arc::new(StringArray::from(vec![nulls; 2])),
+ ])?;
+ let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(
+ expected.into_iter().map(Some),
+ );
+ assert_eq!(result.as_ref(), &expected as &dyn Array, "{nulls}");
+ }
+ Ok(())
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org