This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 59eeb442a6 perf: optimize `rank` for byte views via caching prefix
keys (#10605)
59eeb442a6 is described below
commit 59eeb442a6a8f88823e3cc498623ff1a4eefe928
Author: Thefool <[email protected]>
AuthorDate: Tue Sep 1 15:01:38 2026 +0800
perf: optimize `rank` for byte views via caching prefix keys (#10605)
# Which issue does this PR close?
- Closes #10604.
# Rationale for this change
`byte_view_rank` currently materializes full byte slices and compares
them throughout the sort. For view values longer than the 12-byte inline
capacity, this can repeatedly dereference backing buffers in the O(n log
n) comparison loop.
The benchmark coverage for this optimization is split into #10772 so it
can be present on `main` and be measured by `@adriangbot` without
including benchmark changes in this implementation PR.
# What changes are included in this PR?
- Sort fully inline views using their precomputed `u128` inline keys.
- Cache a big-endian 16-byte key once for long and mixed views,
resolving full values only when keys collide.
- Sample up to eight valid values across two local windows and fall back
to the existing slice path when the sampled key-collision rate is high.
- Return early from sampling when four equal keys already guarantee the
fallback threshold.
- Add `StringView` and `BinaryView` correctness tests for key
collisions, zero bytes, prefix-length relationships, duplicates, nulls,
and all sort option combinations.
# Are these changes tested?
- `cargo test -p arrow-ord rank` passed (10 tests).
- `cargo fmt --all -- --check` passed.
- `git diff --check` passed.
The benchmark cases and the baseline-vs-optimized benchmark run are
tracked in #10772 and will be rerun after that PR is merged.
# Are there any user-facing changes?
No API or behavior changes. This only changes the internal rank
implementation and adds correctness coverage.
---
arrow-ord/src/rank.rs | 384 +++++++++++++++++++++++++++++++++++++++++++++++++-
1 file changed, 380 insertions(+), 4 deletions(-)
diff --git a/arrow-ord/src/rank.rs b/arrow-ord/src/rank.rs
index a4808b93fb..034a002b5d 100644
--- a/arrow-ord/src/rank.rs
+++ b/arrow-ord/src/rank.rs
@@ -106,16 +106,257 @@ fn byte_view_rank<T: ByteViewType>(
array: &GenericByteViewArray<T>,
options: SortOptions,
) -> Vec<u32> {
- let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n|
n.null_count() > 0) {
+ // An inline view already contains the complete value. Convert it once to
+ // a key whose integer ordering matches the byte ordering, as is done by
+ // `sort_byte_view`.
+ if array.data_buffers().is_empty() {
+ let to_sort: Vec<(u128, u32)> = match array.nulls().filter(|n|
n.null_count() > 0) {
+ Some(n) => n
+ .valid_indices()
+ .map(|idx| {
+ // SAFETY: `valid_indices` only yields indices in the
array.
+ let raw = unsafe { *array.views().get_unchecked(idx) };
+ (GenericByteViewArray::<T>::inline_key_fast(raw), idx as
u32)
+ })
+ .collect(),
+ None => array
+ .views()
+ .iter()
+ .enumerate()
+ .map(|(idx, raw)|
(GenericByteViewArray::<T>::inline_key_fast(*raw), idx as u32))
+ .collect(),
+ };
+ return rank_impl(
+ array.len(),
+ to_sort,
+ options,
+ |a, b| a.cmp(&b),
+ |a, b| a == b,
+ );
+ }
+
+ if has_high_byte_view_key_collision_rate(array) {
+ let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n|
n.null_count() > 0) {
+ Some(n) => n
+ .valid_indices()
+ .map(|idx| (array.value(idx).as_ref(), idx as u32))
+ .collect(),
+ None => (0..array.len())
+ .map(|idx| (array.value(idx).as_ref(), idx as u32))
+ .collect(),
+ };
+ return rank_impl(array.len(), to_sort, options, Ord::cmp,
PartialEq::eq);
+ }
+
+ // Cache a wider prefix than the 4 bytes stored in a non-inline view. This
+ // pays for the backing-buffer access once per value instead of once per
+ // comparison, and only resolves the complete value when two keys collide.
+ let to_sort: Vec<(u128, u32)> = match array.nulls().filter(|n|
n.null_count() > 0) {
Some(n) => n
.valid_indices()
- .map(|idx| (array.value(idx).as_ref(), idx as u32))
+ .map(|idx| {
+ // SAFETY: `valid_indices` only yields indices in the array.
+ let value: &[u8] = unsafe {
array.value_unchecked(idx).as_ref() };
+ (byte_view_key(value), idx as u32)
+ })
.collect(),
None => (0..array.len())
- .map(|idx| (array.value(idx).as_ref(), idx as u32))
+ .map(|idx| {
+ // SAFETY: `idx` is in `0..array.len()`.
+ let value: &[u8] = unsafe {
array.value_unchecked(idx).as_ref() };
+ (byte_view_key(value), idx as u32)
+ })
.collect(),
};
- rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq)
+ rank_impl_by(
+ array.len(),
+ to_sort,
+ options,
+ |a, b| compare_view_key(array, a, b),
+ |a, b| equal_view_key(array, a, b),
+ )
+}
+
+// A 16-byte prefix fits in one `u128` and is wider than the 4-byte view
prefix.
+// Shorter keys collide more often; longer keys need another representation.
+const BYTE_VIEW_KEY_LEN: usize = 16;
+
+// Four valid values per window keeps sampling cheap while allowing an early
+// all-collision decision. More samples improve confidence but add buffer
reads.
+const BYTE_VIEW_KEY_SAMPLES_PER_WINDOW: usize = 4;
+
+// Capacity for the two windows; it must be at least
+// `2 * BYTE_VIEW_KEY_SAMPLES_PER_WINDOW`. Increasing it alone has no effect;
+// decreasing it without changing the sampling count can overflow the array.
+const BYTE_VIEW_KEY_SAMPLE_SIZE: usize = 8;
+
+// Bound entries inspected when nulls are present. A higher limit finds valid
+// samples more reliably but costs reads; a lower limit is cheaper but less
+// informative for null-heavy arrays.
+const BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW: usize = 32;
+
+// `colliding_keys * 3 >= sample_len` means roughly one third of sampled keys
+// are duplicates. Lower values fall back earlier; higher values risk keeping
+// the key path for collision-heavy inputs.
+const BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO: usize = 3;
+
+/// Estimates whether cached byte-view keys collide often enough to make the
+/// key-based ranking path unattractive.
+/// Caching a wider key usually avoids repeated backing-buffer reads for long
+/// views. If sampled keys collide frequently, resolving full values plus the
+/// extra key comparison can be slower than comparing slices directly, so the
+/// caller falls back to that path. The bounded two-window sample keeps this
+/// check inexpensive.
+fn has_high_byte_view_key_collision_rate<T: ByteViewType>(array:
&GenericByteViewArray<T>) -> bool {
+ if array.len() < 2 {
+ return false;
+ }
+
+ let mut keys = [0_u128; BYTE_VIEW_KEY_SAMPLE_SIZE];
+ let mut sample_len = 0;
+ let midpoint = array.len() / 2;
+
+ // Probe small local windows in both halves. Keeping each probe local
avoids
+ // turning collision detection itself into scattered backing-buffer reads.
+ for (start, end) in [(0, midpoint), (midpoint, array.len())] {
+ let probe_end =
end.min(start.saturating_add(BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW));
+ let window_start = sample_len;
+ let mut window_samples = 0;
+
+ for idx in start..probe_end {
+ if array.is_null(idx) {
+ continue;
+ }
+
+ // SAFETY: `idx` is within a window bounded by `array.len()`.
+ let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() };
+ keys[sample_len] = byte_view_key(value);
+ sample_len += 1;
+ window_samples += 1;
+ if window_samples == BYTE_VIEW_KEY_SAMPLES_PER_WINDOW {
+ break;
+ }
+ }
+
+ // Four equal keys already contribute three collisions. Even if every
+ // sample in the other window is distinct, that satisfies the final
+ // one-third threshold, so avoid touching the second backing-buffer
+ // window in the common all-collision case.
+ if window_samples == BYTE_VIEW_KEY_SAMPLES_PER_WINDOW
+ && keys[window_start..sample_len]
+ .windows(2)
+ .all(|w| w[0] == w[1])
+ {
+ return true;
+ }
+ }
+
+ if sample_len < 2 {
+ return false;
+ }
+
+ let keys = &mut keys[..sample_len];
+ keys.sort_unstable();
+ let unique_keys = 1 + keys.windows(2).filter(|w| w[0] != w[1]).count();
+
+ // If at least roughly one third of sampled keys collide, comparing the
+ // wider key before every full-value comparison is likely more expensive
+ // than sorting slices directly.
+ let colliding_keys = sample_len - unique_keys;
+ colliding_keys * BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO >= sample_len
+}
+
+#[inline(always)]
+fn byte_view_key(value: &[u8]) -> u128 {
+ let mut key = [0_u8; BYTE_VIEW_KEY_LEN];
+ let key_len = value.len().min(key.len());
+ key[..key_len].copy_from_slice(&value[..key_len]);
+
+ // Big-endian conversion makes integer comparison equivalent to comparing
+ // these bytes lexicographically. Equal keys fall back to the full values,
+ // covering values that differ after 16 bytes and prefixes containing zero.
+ u128::from_be_bytes(key)
+}
+
+#[inline(always)]
+fn compare_view_key<T: ByteViewType>(
+ array: &GenericByteViewArray<T>,
+ a: &(u128, u32),
+ b: &(u128, u32),
+) -> Ordering {
+ match a.0.cmp(&b.0) {
+ Ordering::Equal => {
+ // SAFETY: both indices were produced from this array above.
+ let full_a: &[u8] = unsafe { array.value_unchecked(a.1 as
usize).as_ref() };
+ let full_b: &[u8] = unsafe { array.value_unchecked(b.1 as
usize).as_ref() };
+ full_a.cmp(full_b)
+ }
+ ordering => ordering,
+ }
+}
+
+#[inline(always)]
+fn equal_view_key<T: ByteViewType>(
+ array: &GenericByteViewArray<T>,
+ a: &(u128, u32),
+ b: &(u128, u32),
+) -> bool {
+ if a.0 != b.0 {
+ return false;
+ }
+
+ // SAFETY: both indices were produced from this array above.
+ let full_a: &[u8] = unsafe { array.value_unchecked(a.1 as usize).as_ref()
};
+ let full_b: &[u8] = unsafe { array.value_unchecked(b.1 as usize).as_ref()
};
+ full_a == full_b
+}
+
+fn rank_impl_by<T, C, E>(
+ len: usize,
+ mut valid: Vec<(T, u32)>,
+ options: SortOptions,
+ compare: C,
+ eq: E,
+) -> Vec<u32>
+where
+ C: Fn(&(T, u32), &(T, u32)) -> Ordering,
+ E: Fn(&(T, u32), &(T, u32)) -> bool,
+{
+ // Same ranking and null handling as `rank_impl`, but callbacks receive
+ // tuple references because key collisions use the index to read the full
+ // value. `rank_impl` passes copied values directly through `Fn(T, T)`.
+ // We can use an unstable sort as we combine equal values later
+ valid.sort_unstable_by(compare);
+ if options.descending {
+ valid.reverse();
+ }
+
+ let (mut valid_rank, null_rank) = match options.nulls_first {
+ true => (len as u32, (len - valid.len()) as u32),
+ false => (valid.len() as u32, len as u32),
+ };
+
+ let mut out: Vec<_> = vec![null_rank; len];
+ if let Some(v) = valid.last() {
+ out[v.1 as usize] = valid_rank;
+ }
+
+ let mut count = 1; // Number of values in rank
+ for w in valid.windows(2).rev() {
+ match eq(&w[0], &w[1]) {
+ true => {
+ count += 1;
+ out[w[0].1 as usize] = valid_rank;
+ }
+ false => {
+ valid_rank -= count;
+ count = 1;
+ out[w[0].1 as usize] = valid_rank
+ }
+ }
+ }
+
+ out
}
fn rank_impl<T, C, E>(
@@ -247,6 +488,21 @@ mod tests {
use super::*;
use arrow_array::*;
+ fn assert_same_rank(left: &dyn Array, right: &dyn Array) {
+ for descending in [false, true] {
+ for nulls_first in [false, true] {
+ let options = SortOptions {
+ descending,
+ nulls_first,
+ };
+ assert_eq!(
+ rank(left, Some(options)).unwrap(),
+ rank(right, Some(options)).unwrap()
+ );
+ }
+ }
+ }
+
#[test]
fn test_primitive() {
let descending = SortOptions {
@@ -389,6 +645,35 @@ mod tests {
assert_eq!(res, &[3, 1, 4, 3]);
}
+ #[test]
+ fn test_inline_byte_views() {
+ let string_values = vec![
+ Some(""),
+ Some("short"),
+ None,
+ Some("0123456789qa"), // exactly the 12-byte inline limit
+ Some("short"),
+ ];
+ let string_view = StringViewArray::from(string_values.clone());
+ let string = StringArray::from(string_values);
+
+ assert!(string_view.data_buffers().is_empty());
+ assert_same_rank(&string_view, &string);
+
+ let binary_values: Vec<Option<&[u8]>> = vec![
+ Some(b""),
+ Some(b"short"),
+ None,
+ Some(b"0123456789qa"),
+ Some(b"short"),
+ ];
+ let binary_view = BinaryViewArray::from_iter(binary_values.clone());
+ let binary = BinaryArray::from_opt_vec(binary_values);
+
+ assert!(binary_view.data_buffers().is_empty());
+ assert_same_rank(&binary_view, &binary);
+ }
+
#[test]
fn test_string_view_with_nulls() {
let values = StringViewArray::from(vec![
@@ -413,4 +698,95 @@ mod tests {
let res = rank(&values, None).unwrap();
assert_eq!(res, &[3, 4, 1, 3]);
}
+
+ #[test]
+ fn test_string_view_key_collisions() {
+ let values = vec![
+ Some("abcdefghijklmnop"),
+ Some("abcdefghijklmnopA"),
+ Some("abcdefghijklmnopB"),
+ Some("abcdefghijklmnopA"),
+ Some("abcdefghijklmno"),
+ Some("short"),
+ None,
+ ];
+ let expected = StringArray::from(values.clone());
+ let actual = StringViewArray::from(values);
+
+ assert_same_rank(&actual, &expected);
+ }
+
+ #[test]
+ fn test_binary_view_key_collisions() {
+ let zeroes_16 = [0_u8; 16];
+ let zeroes_17 = [0_u8; 17];
+ let mut zeroes_then_one = [0_u8; 17];
+ zeroes_then_one[16] = 1;
+ let mut zeroes_then_two = [0_u8; 17];
+ zeroes_then_two[16] = 2;
+
+ let values: Vec<Option<&[u8]>> = vec![
+ Some(b""),
+ Some(b"\0"),
+ Some(&zeroes_16),
+ Some(&zeroes_17),
+ Some(&zeroes_then_one),
+ Some(&zeroes_then_two),
+ Some(&zeroes_then_one),
+ None,
+ ];
+ let expected = BinaryArray::from_opt_vec(values.clone());
+ let actual = BinaryViewArray::from_iter(values);
+
+ assert_same_rank(&actual, &expected);
+ }
+
+ #[test]
+ fn test_byte_view_high_key_collision_detection() {
+ const SIZE: u32 = 64;
+
+ let same_key: StringViewArray = (0..SIZE)
+ .map(|i| {
+ let suffix = i.wrapping_mul(2_654_435_761);
+ Some(format!("abcdefghijklmnop{suffix:08x}"))
+ })
+ .collect();
+
+ assert_eq!(
+ byte_view_key(same_key.value(0).as_bytes()),
+ byte_view_key(same_key.value(1).as_bytes())
+ );
+ assert!(has_high_byte_view_key_collision_rate(&same_key));
+
+ let clustered: StringViewArray = (0..SIZE)
+ .map(|i| {
+ let suffix = i.wrapping_mul(2_654_435_761);
+ let value = if i < SIZE / 2 {
+ format!("{suffix:016x}abcdefgh")
+ } else {
+ format!("abcdefghijklmnop{suffix:08x}")
+ };
+ Some(value)
+ })
+ .collect();
+ assert!(has_high_byte_view_key_collision_rate(&clustered));
+
+ let with_nulls: StringViewArray = (0..SIZE)
+ .map(|i| {
+ (i % 2 == 0).then(|| {
+ let suffix = i.wrapping_mul(2_654_435_761);
+ format!("abcdefghijklmnop{suffix:08x}")
+ })
+ })
+ .collect();
+ assert!(has_high_byte_view_key_collision_rate(&with_nulls));
+
+ let distinct: StringViewArray = (0..SIZE)
+ .map(|i| {
+ let suffix = i.wrapping_mul(2_654_435_761);
+ Some(format!("{suffix:016x}abcdefgh"))
+ })
+ .collect();
+ assert!(!has_high_byte_view_key_collision_rate(&distinct));
+ }
}
