Skip to content

Sort performance #9407

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cetra3 wants to merge 1 commit into
base: main
Choose a base branch
from
+59 −56
Open

Sort performance #9407

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 59 additions & 56 deletions arrow-ord/src/sort.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,55 +74,65 @@ where
{
let sort_options = options.unwrap_or_default();

let mut mutable_buffer = vec![T::default_value(); primitive_values.len()];
let mutable_slice = &mut mutable_buffer;

let input_values = primitive_values.values().as_ref();

let nulls_count = primitive_values.null_count();
let valid_count = primitive_values.len() - nulls_count;
if let Some(nulls) = primitive_values.nulls().filter(|n| n.null_count() > 0) {
let length = primitive_values.len();
let nulls_count = nulls.null_count();
let valid_count = length - nulls_count;

let mut validity_buffer = BooleanBufferBuilder::new(length);
let mut mutable_buffer = Vec::with_capacity(length);

let values_slice = if sort_options.nulls_first {
validity_buffer.append_n(nulls_count, false);
validity_buffer.append_n(valid_count, true);

mutable_buffer.resize(nulls_count, T::default_value());
// SAFETY: The indices are valid because they come from the nulls buffer.
mutable_buffer.extend(
nulls
.valid_indices()
.map(|i| unsafe { primitive_values.value_unchecked(i) }),
);

let null_bit_buffer = match nulls_count > 0 {
true => {
let mut validity_buffer = BooleanBufferBuilder::new(primitive_values.len());
if sort_options.nulls_first {
validity_buffer.append_n(nulls_count, false);
validity_buffer.append_n(valid_count, true);
} else {
validity_buffer.append_n(valid_count, true);
validity_buffer.append_n(nulls_count, false);
}
Some(validity_buffer.finish().into())
}
false => None,
};
&mut mutable_buffer[nulls_count..]
} else {
validity_buffer.append_n(valid_count, true);
validity_buffer.append_n(nulls_count, false);

// SAFETY: The indices are valid because they come from the nulls buffer.
mutable_buffer.extend(
nulls
.valid_indices()
.map(|i| unsafe { primitive_values.value_unchecked(i) }),
);
mutable_buffer.resize(length, T::default_value());

if let Some(nulls) = primitive_values.nulls().filter(|n| n.null_count() > 0) {
let values_slice = match sort_options.nulls_first {
true => &mut mutable_slice[nulls_count..],
false => &mut mutable_slice[..valid_count],
&mut mutable_buffer[..valid_count]
};

for (write_index, index) in nulls.valid_indices().enumerate() {
values_slice[write_index] = primitive_values.value(index);
}

values_slice.sort_unstable_by(|a, b| a.compare(*b));
if sort_options.descending {
values_slice.reverse();
values_slice.sort_unstable_by(|a, b| b.compare(*a));
} else {
values_slice.sort_unstable_by(|a, b| a.compare(*b));
}

return Ok(Arc::new(
PrimitiveArray::<T>::try_new(mutable_buffer.into(), Some(validity_buffer.into()))?
.with_data_type(primitive_values.data_type().clone()),
));
} else {
mutable_slice.copy_from_slice(input_values);
mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
let mut mutable_buffer = primitive_values.values().to_vec();
if sort_options.descending {
mutable_slice.reverse();
mutable_buffer.sort_unstable_by(|a, b| b.compare(*a));
} else {
mutable_buffer.sort_unstable_by(|a, b| a.compare(*b));
}
}

Ok(Arc::new(
PrimitiveArray::<T>::try_new(mutable_buffer.into(), null_bit_buffer)?
.with_data_type(primitive_values.data_type().clone()),
))
return Ok(Arc::new(
PrimitiveArray::<T>::try_new(mutable_buffer.into(), None)?
.with_data_type(primitive_values.data_type().clone()),
));
}
}

/// Sort the `ArrayRef` partially.
Expand Down Expand Up @@ -216,27 +226,20 @@ fn partition_validity_scan(
let mut valid = Vec::with_capacity(len - null_count);
let mut nulls = Vec::with_capacity(null_count);

unsafe {
// 1) Write valid indices (bits == 1)
let valid_slice = valid.spare_capacity_mut();
for (i, idx) in bitmap.inner().set_indices_u32().enumerate() {
valid_slice[i].write(idx);
}
let mut cur_idx = 0u32;

// 2) Write null indices by inverting
let inv_buf = !bitmap.inner();
Copy link
Copy Markdown
Contributor Author

@cetra3 cetra3 Feb 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is where most of the saving comes in. Recreating the bitmap with inverse is another pass over it, so we have to read it twice.

Copy link
Copy Markdown
Contributor

@Dandandan Dandandan Feb 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense

Copy link
Copy Markdown
Contributor

@jhorstmann jhorstmann Feb 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting, I think this was changed in the other direction in 625e6ee#diff-2c3ee46148ccba8c1092294caab66ba73416314595ce2eef9b2bdd9cd898b27d

I wonder what changed now that makes the benchmarks go the other way.

let null_slice = nulls.spare_capacity_mut();
for (i, idx) in inv_buf.set_indices_u32().enumerate() {
null_slice[i].write(idx);
bitmap.inner().iter().for_each(|is_valid| {
if is_valid {
valid.push(cur_idx);
} else {
nulls.push(cur_idx);
}

// Finalize lengths
valid.set_len(len - null_count);
nulls.set_len(null_count);
}
cur_idx += 1;
});

assert_eq!(valid.len(), len - null_count);
assert_eq!(nulls.len(), null_count);
debug_assert_eq!(valid.len(), len - null_count);
debug_assert_eq!(nulls.len(), null_count);
(valid, nulls)
}

Expand Down
Loading