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/09/01 15:00:57 UTC
(datafusion) branch gh-readonly-queue/main/pr-24823-80ae9e51c289ef3b9470a5ff989118a391f390df created (now da89c7c85b)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24823-80ae9e51c289ef3b9470a5ff989118a391f390df
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at da89c7c85b array_sort / array_has: monomorphize per native type instead of per Arrow type (reduce binary size) (#24823)
This branch includes the following new commits:
new da89c7c85b array_sort / array_has: monomorphize per native type instead of per Arrow type (reduce binary size) (#24823)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: array_sort / array_has: monomorphize per native type instead of per Arrow type (reduce binary size) (#24823)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24823-80ae9e51c289ef3b9470a5ff989118a391f390df
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit da89c7c85b1c888ba4038317eacbdcfe03ab9b1a
Author: Jay Zhan <ja...@gmail.com>
AuthorDate: Tue Sep 1 15:00:34 2026 +0000
array_sort / array_has: monomorphize per native type instead of per Arrow type (reduce binary size) (#24823)
## Which issue does this PR close?
- Part of https://github.com/apache/datafusion/issues/24727
## Rationale for this change
`array_sort`, `array_has` and `array_min`/`array_max` each dispatch to a
kernel
that is generic over **both** the Arrow element type and the list offset
width,
so each kernel is monomorphized ~64 times. The bodies only ever move and
compare
native values, so most of those copies are byte-identical: `Int32`,
`Date32`,
`Time32Second` and `IntervalYearMonth` all sort and compare as `i32`.
`array_sort` is the worst case, and pays for it twice over: it calls
`sort_unstable_by` with a separate ascending and descending comparator
closure,
and each distinct closure type drags in its own full copy of the
standard
library's `ipnsort`. That alone accounts for **437K IR lines across
7,832
instantiations — 30% of the entire crate**.
## What changes are included in this PR?
- `sort.rs`: the per-row sort kernels become generic over the native
value type
rather than the Arrow type, and both sort directions route through a
single
comparator. `ArrowNativeTypeOp::compare` is a total order, so a
descending
sort is exactly the reverse of the ascending one — sorting one way and
reversing keeps one instantiation of the standard library's sort per
native
type instead of one per (Arrow type, offset width, direction) triple.
- `sort.rs`: the row-index sort in the non-primitive path moves into a
helper
free of the offset width, so it is instantiated once rather than once
per
list type.
- `array_has.rs`: the primitive scan is split into a thin per-Arrow-type
shim
and a native-typed kernel, so the equality scan is compiled once per
native
type.
- `min_max.rs`: `scalar_min_max` becomes generic over the native type.
The kernels stay generic over `OffsetSizeTrait` deliberately. Erasing
that
dimension (widening offsets to `&[usize]`) was measured and reverted: it
costs
an allocation and a pass over every row, which regressed `array_sort`
over
5-element lists by ~30%, and it saves only ~15K IR lines because the
expensive
instantiation is keyed on the comparator closure, which never mentions
the
offset width. That is also why `min_max` improves only slightly here —
erasing
its offset dimension needs a zero-allocation offsets view (an enum over
`&[i32]`/`&[i64]`) to be worth doing, which is left as a follow-up.
## Code size
### 1. Generated code: `cargo llvm-lines --release -p
datafusion-functions-nested --lib`
| | main | this PR | change |
| --- | ---: | ---: | ---: |
| **crate total (lines)** | 1,455,442 | 965,772 | **−489,670 (−33.6%)**
|
| **crate total (copies)** | 26,832 | 18,219 | −8,613 (−32.1%) |
| `sort` module | 537,199 | 75,615 | −461,584 (−85.9%) |
| `array_has` module | 73,292 | 46,948 | −26,344 (−35.9%) |
| `min_max` module | 26,316 | 24,016 | −2,300 (−8.7%) |
The standard library's sort machinery (`core::slice::sort`, attributed
to this
crate's comparators) goes from **437,623 IR lines in 7,832
instantiations to
27,851 in 579**.
### 2. Final binary size: release bench binaries
| binary | main | this PR | change |
| --- | ---: | ---: | ---: |
| `array_sort` | 6,478,112 | 6,213,728 | −264,384 (−4.1%) |
| `array_has` | 6,825,968 | 6,693,872 | −132,096 (−1.9%) |
| `array_min_max` | 7,536,608 | 7,536,608 | 0 |
## Are these changes tested?
Yes, by existing tests — this is a refactor with no behavior change:
- `cargo test -p datafusion-functions-nested --lib` — 112 tests pass
- `cargo test -p datafusion-sqllogictest --test sqllogictests -- array`
— all 53
files pass
- `cargo clippy --all-targets --all-features -- -D warnings` clean
### Benchmarks
`array_sort`, `array_min_max` and `array_has`, run as base/branch pairs
from
pre-built binaries so no compilation interleaves with measurement,
repeated
across passes with a base-vs-base noise floor measured alongside.
No benchmark is slower beyond the noise floor. Two are consistently
faster:
| benchmark | change |
| --- | ---: |
| `array_has_array_i64/found/10` | −17% |
| `array_has_array_null_patterns/i64/nulls30_found` | −6% |
Caveat on the numbers: these were measured on a laptop whose best
observed
noise floor was ~2%, so differences under that are not resolvable.
Everything
outside those two rows landed within ±2%.
## Are there any user-facing changes?
No. No public API or behavior changes; `array_sort` output ordering is
unchanged, including for floats (`compare` is a total order, so
reversing the
ascending order reproduces the descending order exactly, `NaN` and
`-0.0`
included).
---
datafusion/functions-nested/benches/array_sort.rs | 50 +++++-
datafusion/functions-nested/src/array_has.rs | 45 +++---
datafusion/functions-nested/src/min_max.rs | 13 +-
datafusion/functions-nested/src/sort.rs | 183 +++++++++++-----------
4 files changed, 168 insertions(+), 123 deletions(-)
diff --git a/datafusion/functions-nested/benches/array_sort.rs b/datafusion/functions-nested/benches/array_sort.rs
index 940c0396cb..2ceee91f4c 100644
--- a/datafusion/functions-nested/benches/array_sort.rs
+++ b/datafusion/functions-nested/benches/array_sort.rs
@@ -22,6 +22,7 @@ use arrow::array::{ArrayRef, BooleanBufferBuilder, Int32Array, ListArray, String
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::datatypes::{DataType, Field};
use criterion::{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;
@@ -119,9 +120,25 @@ fn create_string_list_array(num_rows: usize, elements_per_row: usize) -> ArrayRe
}
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(
+ udf: &ArraySort,
+ array: &ArrayRef,
+ order: Option<&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());
+ }
udf.invoke_with_args(ScalarFunctionArgs {
- args: vec![ColumnarValue::Array(Arc::clone(array))],
- arg_fields: vec![Field::new("arr", array.data_type().clone(), true).into()],
+ args,
+ arg_fields,
number_rows: array.len(),
return_field: Field::new("result", array.data_type().clone(), true).into(),
config_options: Arc::new(ConfigOptions::default()),
@@ -149,6 +166,21 @@ fn bench_array_sort(c: &mut Criterion) {
);
}
+ // 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(
+ BenchmarkId::new("int32_desc", elements_per_row),
+ &elements_per_row,
+ |b, _| {
+ b.iter(|| {
+ black_box(invoke_array_sort_with_order(&udf, &array, Some("DESC")));
+ });
+ },
+ );
+ }
+
// Int32 with nulls in the outer list (10% null rows), single size
{
let array = create_int32_list_array(NUM_ROWS, 50, true);
@@ -174,6 +206,20 @@ fn bench_array_sort(c: &mut Criterion) {
);
}
+ 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")));
+ });
+ },
+ );
+ }
+
// String arrays
for &elements_per_row in &[5, 20, 100, 1000] {
let array = create_string_list_array(NUM_ROWS, elements_per_row);
diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs
index 24a9913d2c..95023aedd6 100644
--- a/datafusion/functions-nested/src/array_has.rs
+++ b/datafusion/functions-nested/src/array_has.rs
@@ -18,9 +18,8 @@
//! [`ScalarUDFImpl`] definitions for array_has, array_has_all and array_has_any functions.
use arrow::array::{
- Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray,
- BooleanBufferBuilder, Datum, MAX_INLINE_VIEW_LEN, PrimitiveArray, Scalar,
- StringArrayType, StringViewArray,
+ Array, ArrayRef, ArrowNativeTypeOp, AsArray, BooleanArray, BooleanBufferBuilder,
+ Datum, MAX_INLINE_VIEW_LEN, Scalar, StringArrayType, StringViewArray,
};
use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
use arrow::datatypes::DataType;
@@ -359,7 +358,7 @@ fn array_has_dispatch_for_array(
None
} else {
downcast_primitive_array! {
- visible_values => {
+ visible_values, needle => {
// The element-null path makes several passes over the values, so
// past a large average list length the per-row `eq` kernel is
// faster -- bail to it. The single-pass all-valid path has no such
@@ -371,25 +370,28 @@ fn array_has_dispatch_for_array(
{
None
} else {
- Some(array_has_array_primitive(
- visible_values, needle, &offsets,
+ Some(array_has_array_native(
+ visible_values.values(),
+ visible_values.nulls(),
+ needle.values(),
+ &offsets,
combined_nulls.as_ref(),
))
}
- },
- DataType::Utf8 => Some(array_has_array_string(
+ }
+ (DataType::Utf8, _) => Some(array_has_array_string(
visible_values.as_string::<i32>(),
needle.as_string::<i32>(),
&offsets,
combined_nulls.as_ref(),
)),
- DataType::LargeUtf8 => Some(array_has_array_string(
+ (DataType::LargeUtf8, _) => Some(array_has_array_string(
visible_values.as_string::<i64>(),
needle.as_string::<i64>(),
&offsets,
combined_nulls.as_ref(),
)),
- DataType::Utf8View => Some(array_has_array_string_view(
+ (DataType::Utf8View, _) => Some(array_has_array_string_view(
visible_values.as_string_view(),
needle.as_string_view(),
&offsets,
@@ -430,21 +432,20 @@ const NULL_FAST_PATH_MAX_LEN: usize = 512;
/// 2. Nulls: AND the equality bitmap with validity (a null slot's value is
/// arbitrary), then reduce each row to "any bit set". Chunked to bound the
/// expanded needle.
-fn array_has_array_primitive<T: ArrowPrimitiveType>(
- values: &PrimitiveArray<T>,
- needle: &dyn Array,
+///
+/// Generic over the native type alone -- the caller's `downcast_primitive_array!`
+/// peels off the Arrow type -- so `Int32`, `Date32` and `Time32` share a single
+/// instantiation instead of getting one each.
+fn array_has_array_native<N: ArrowNativeTypeOp>(
+ value_slice: &[N],
+ element_nulls: Option<&NullBuffer>,
+ needle_slice: &[N],
offsets: &[usize],
combined_nulls: Option<&NullBuffer>,
-) -> BooleanBuffer
-where
- T::Native: ArrowNativeTypeOp,
-{
- let needle = needle.as_primitive::<T>();
+) -> BooleanBuffer {
let num_rows = offsets.len() - 1;
- let value_slice = values.values();
- let needle_slice = needle.values();
- let Some(element_nulls) = values.nulls() else {
+ let Some(element_nulls) = element_nulls else {
return BooleanBuffer::collect_bool(num_rows, |i| {
if combined_nulls.is_some_and(|n| n.is_null(i)) {
return false;
@@ -461,7 +462,7 @@ where
// Case 2 (see fn doc), chunked like the all/any kernels.
let mut result = BooleanBufferBuilder::new(num_rows);
- let mut needle_expanded: Vec<T::Native> = Vec::new();
+ let mut needle_expanded: Vec<N> = Vec::new();
for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) {
let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows);
let elem_start = offsets[chunk_start];
diff --git a/datafusion/functions-nested/src/min_max.rs b/datafusion/functions-nested/src/min_max.rs
index ba9e5e7a07..2e7a32ba31 100644
--- a/datafusion/functions-nested/src/min_max.rs
+++ b/datafusion/functions-nested/src/min_max.rs
@@ -21,6 +21,7 @@ use arrow::array::{
Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, GenericListArray,
OffsetSizeTrait, PrimitiveBuilder, downcast_primitive,
};
+use arrow::buffer::NullBuffer;
use arrow::datatypes::DataType;
use arrow::datatypes::DataType::{LargeList, List};
use datafusion_common::Result;
@@ -263,7 +264,7 @@ fn primitive_array_min_max<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
match len {
0 => None,
_ if len < ARROW_COMPUTE_THRESHOLD => {
- scalar_min_max::<T>(values_slice, values_nulls, start, end, is_min)
+ scalar_min_max(values_slice, values_nulls, start, end, is_min)
}
_ => {
let slice = values_array.slice(start, len);
@@ -285,14 +286,14 @@ fn primitive_array_min_max<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
/// Computes min or max for a single list row by directly scanning a slice of
/// the flat values buffer.
#[inline]
-fn scalar_min_max<T: ArrowPrimitiveType>(
- values_slice: &[T::Native],
- values_nulls: Option<&arrow::buffer::NullBuffer>,
+fn scalar_min_max<N: ArrowNativeTypeOp>(
+ values_slice: &[N],
+ values_nulls: Option<&NullBuffer>,
start: usize,
end: usize,
is_min: bool,
-) -> Option<T::Native> {
- let mut best: Option<T::Native> = None;
+) -> Option<N> {
+ let mut best: Option<N> = None;
for (i, &val) in values_slice[start..end].iter().enumerate() {
if let Some(nulls) = values_nulls
&& !nulls.is_valid(start + i)
diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs
index f4f9148f76..fadc5da4b4 100644
--- a/datafusion/functions-nested/src/sort.rs
+++ b/datafusion/functions-nested/src/sort.rs
@@ -25,11 +25,11 @@ use arrow::array::{
};
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::datatypes::{ArrowNativeTypeOp, DataType, FieldRef};
-use arrow::row::{RowConverter, SortField};
+use arrow::row::{RowConverter, Rows, SortField};
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;
-use datafusion_common::{Result, exec_err, internal_datafusion_err};
+use datafusion_common::{Result, exec_err};
use datafusion_expr::{
ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation,
ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
@@ -238,99 +238,104 @@ fn sort_primitive_list<T: ArrowPrimitiveType, OffsetSize: OffsetSizeTrait>(
where
T::Native: ArrowNativeTypeOp,
{
- if prim_values.null_count() > 0 {
- sort_list_with_nulls(prim_values, list_array, field, sort_options)
- } else {
- sort_list_no_nulls(prim_values, list_array, field, sort_options)
+ let descending = sort_options.is_some_and(|o| o.descending);
+ let nulls_first = sort_options.is_none_or(|o| o.nulls_first);
+ let list_nulls = list_array.nulls();
+ let offsets = list_array.offsets();
+
+ let (values, validity) = match prim_values.nulls() {
+ Some(element_nulls) if element_nulls.null_count() > 0 => {
+ let (values, validity) = sort_rows_with_nulls(
+ prim_values.values(),
+ element_nulls,
+ offsets,
+ list_nulls,
+ descending,
+ nulls_first,
+ );
+ (values, Some(validity))
+ }
+ _ => (
+ sort_rows_no_nulls(prim_values.values(), offsets, list_nulls, descending),
+ None,
+ ),
+ };
+
+ let sorted_values = Arc::new(
+ PrimitiveArray::<T>::new(values.into(), validity)
+ .with_data_type(prim_values.data_type().clone()),
+ );
+
+ Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
+ field,
+ rebase_offsets(offsets),
+ sorted_values,
+ list_nulls.cloned(),
+ )?))
+}
+
+/// Sorts one row's values in place.
+///
+/// [`ArrowNativeTypeOp::compare`] is a total order, so a descending sort is the
+/// reverse of the ascending one. Sorting one way and reversing keeps a single
+/// instantiation of the standard library's sort per native type, instead of one
+/// per comparator closure.
+#[inline]
+fn sort_row<N: ArrowNativeTypeOp>(row: &mut [N], descending: bool) {
+ row.sort_unstable_by(|a, b| a.compare(*b));
+ if descending {
+ row.reverse();
}
}
/// Fast path for primitive values with no element-level nulls. Copies all
/// values into a single `Vec` and sorts each row's slice in-place.
-fn sort_list_no_nulls<T: ArrowPrimitiveType, OffsetSize: OffsetSizeTrait>(
- prim_values: &PrimitiveArray<T>,
- list_array: &GenericListArray<OffsetSize>,
- field: FieldRef,
- sort_options: Option<SortOptions>,
-) -> Result<ArrayRef>
-where
- T::Native: ArrowNativeTypeOp,
-{
- let row_count = list_array.len();
- let offsets = list_array.offsets();
+fn sort_rows_no_nulls<N: ArrowNativeTypeOp, O: OffsetSizeTrait>(
+ src_values: &[N],
+ offsets: &[O],
+ list_nulls: Option<&NullBuffer>,
+ descending: bool,
+) -> Vec<N> {
let values_start = offsets[0].as_usize();
- let values_end = offsets[row_count].as_usize();
-
- let descending = sort_options.is_some_and(|o| o.descending);
+ let values_end = offsets[offsets.len() - 1].as_usize();
// Copy all values into a mutable buffer
- let mut values: Vec<T::Native> =
- prim_values.values()[values_start..values_end].to_vec();
+ let mut values = src_values[values_start..values_end].to_vec();
for (row_index, window) in offsets.windows(2).enumerate() {
- if list_array.is_null(row_index) {
+ if list_nulls.is_some_and(|n| n.is_null(row_index)) {
continue;
}
let start = window[0].as_usize() - values_start;
let end = window[1].as_usize() - values_start;
- let slice = &mut values[start..end];
- if descending {
- slice.sort_unstable_by(|a, b| b.compare(*a));
- } else {
- slice.sort_unstable_by(|a, b| a.compare(*b));
- }
+ sort_row(&mut values[start..end], descending);
}
- let new_offsets = rebase_offsets(offsets);
- let sorted_values = Arc::new(
- PrimitiveArray::<T>::new(values.into(), None)
- .with_data_type(prim_values.data_type().clone()),
- );
-
- Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
- field,
- new_offsets,
- sorted_values,
- list_array.nulls().cloned(),
- )?))
+ values
}
/// Slow path for primitive values with element-level nulls.
-fn sort_list_with_nulls<T: ArrowPrimitiveType, OffsetSize: OffsetSizeTrait>(
- prim_values: &PrimitiveArray<T>,
- list_array: &GenericListArray<OffsetSize>,
- field: FieldRef,
- sort_options: Option<SortOptions>,
-) -> Result<ArrayRef>
-where
- T::Native: ArrowNativeTypeOp,
-{
- let row_count = list_array.len();
- let offsets = list_array.offsets();
+fn sort_rows_with_nulls<N: ArrowNativeTypeOp, O: OffsetSizeTrait>(
+ src_values: &[N],
+ src_nulls: &NullBuffer,
+ offsets: &[O],
+ list_nulls: Option<&NullBuffer>,
+ descending: bool,
+ nulls_first: bool,
+) -> (Vec<N>, NullBuffer) {
let values_start = offsets[0].as_usize();
- let values_end = offsets[row_count].as_usize();
- let total_values = values_end - values_start;
+ let total_values = offsets[offsets.len() - 1].as_usize() - values_start;
- let descending = sort_options.is_some_and(|o| o.descending);
- let nulls_first = sort_options.is_none_or(|o| o.nulls_first);
-
- let mut out_values: Vec<T::Native> = vec![T::Native::default(); total_values];
+ let mut out_values: Vec<N> = vec![N::default(); total_values];
let mut validity = BooleanBufferBuilder::new(total_values);
- let src_nulls = prim_values.nulls().ok_or_else(|| {
- internal_datafusion_err!(
- "sort_list_with_nulls called but values have no null buffer"
- )
- })?;
- let src_values = prim_values.values();
-
for (row_index, window) in offsets.windows(2).enumerate() {
let start = window[0].as_usize();
let end = window[1].as_usize();
let row_len = end - start;
let out_start = start - values_start;
- if list_array.is_null(row_index) || row_len == 0 {
+ if list_nulls.is_some_and(|n| n.is_null(row_index)) || row_len == 0 {
validity.append_n(row_len, false);
continue;
}
@@ -342,20 +347,18 @@ where
// 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 in start..end {
- if src_nulls.is_valid(i) {
- out_values[write_pos] = src_values[i];
+ 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;
}
}
- let valid_slice = &mut out_values
- [out_start + valid_offset..out_start + valid_offset + valid_count];
- if descending {
- valid_slice.sort_unstable_by(|a, b| b.compare(*a));
- } else {
- valid_slice.sort_unstable_by(|a, b| a.compare(*b));
- }
+ sort_row(
+ &mut out_values
+ [out_start + valid_offset..out_start + valid_offset + valid_count],
+ descending,
+ );
// Build validity bits
if nulls_first {
@@ -367,20 +370,7 @@ where
}
}
- let new_offsets = rebase_offsets(offsets);
-
- let null_buffer = NullBuffer::from(validity.finish());
- let sorted_values = Arc::new(
- PrimitiveArray::<T>::new(out_values.into(), Some(null_buffer))
- .with_data_type(prim_values.data_type().clone()),
- );
-
- Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
- field,
- new_offsets,
- sorted_values,
- list_array.nulls().cloned(),
- )?))
+ (out_values, NullBuffer::from(validity.finish()))
}
/// Sort a non-pritive-typed ListArray by converting all rows at once using
@@ -426,9 +416,7 @@ fn array_sort_non_primitive<OffsetSize: OffsetSizeTrait>(
if len <= 1 {
indices.extend((local_start..local_start + len).map(OffsetSize::usize_as));
} else {
- sort_scratch.clear();
- sort_scratch.extend(local_start..local_start + len);
- sort_scratch.sort_unstable_by(|&a, &b| rows.row(a).cmp(&rows.row(b)));
+ sort_row_indices(&mut sort_scratch, &rows, local_start, len);
indices.extend(sort_scratch.iter().map(|&i| OffsetSize::usize_as(i)));
}
@@ -449,6 +437,15 @@ fn array_sort_non_primitive<OffsetSize: OffsetSizeTrait>(
)?))
}
+/// Sorts `scratch` to hold the indices `start..start + len` ordered by their
+/// encoded rows. Kept free of the offset width so the standard library's sort
+/// is instantiated once rather than once per list type.
+fn sort_row_indices(scratch: &mut Vec<usize>, rows: &Rows, start: usize, len: usize) {
+ scratch.clear();
+ scratch.extend(start..start + len);
+ scratch.sort_unstable_by(|&a, &b| rows.row(a).cmp(&rows.row(b)));
+}
+
/// Select elements from `values` at the given `indices` using `compute::take`.
/// We consume `indices` in order to avoid an intermediate copy.
fn take_by_indices<OffsetSize: OffsetSizeTrait>(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org