Skip to content

Commit aa96097

Browse files
authored
Perf: Add prefix compare for inlined compare and change use of inline_value to inline it to a u128 (#7748)
# Which issue does this PR close? Closes [#7743](#7743) # Rationale for this change Change the fast path to use u128 to compare for lt case, also for inline <12 case to use u128 to compare. Also when we have > 12 data buffer case, we change 4 bytes compare from each byte compare to u32 compare. # What changes are included in this PR? Change the fast path to use u128 to compare for lt case, also for inline <12 case to use u128 to compare. Also when we have > 12 data buffer case, we change 4 bytes compare from each byte compare to u32 compare. # Are there any user-facing changes? No
1 parent 674dc17 commit aa96097

File tree

2 files changed

+156
-24
lines changed

2 files changed

+156
-24
lines changed

arrow-array/src/array/byte_view_array.rs

Lines changed: 146 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use arrow_schema::{ArrowError, DataType};
2727
use core::str;
2828
use num::ToPrimitive;
2929
use std::any::Any;
30+
use std::cmp::Ordering;
3031
use std::fmt::Debug;
3132
use std::marker::PhantomData;
3233
use std::sync::Arc;
@@ -539,25 +540,30 @@ impl<T: ByteViewType + ?Sized> GenericByteViewArray<T> {
539540
left_idx: usize,
540541
right: &GenericByteViewArray<T>,
541542
right_idx: usize,
542-
) -> std::cmp::Ordering {
543+
) -> Ordering {
543544
let l_view = left.views().get_unchecked(left_idx);
544-
let l_len = *l_view as u32;
545+
let l_byte_view = ByteView::from(*l_view);
545546

546547
let r_view = right.views().get_unchecked(right_idx);
547-
let r_len = *r_view as u32;
548+
let r_byte_view = ByteView::from(*r_view);
548549

549-
if l_len <= MAX_INLINE_VIEW_LEN && r_len <= MAX_INLINE_VIEW_LEN {
550-
let l_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, l_len as usize) };
551-
let r_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, r_len as usize) };
552-
return l_data.cmp(r_data);
550+
let l_len = l_byte_view.length;
551+
let r_len = r_byte_view.length;
552+
553+
if l_len <= 12 && r_len <= 12 {
554+
return Self::inline_key_fast(*l_view).cmp(&Self::inline_key_fast(*r_view));
553555
}
554556

555557
// one of the string is larger than 12 bytes,
556558
// we then try to compare the inlined data first
557-
let l_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, 4) };
558-
let r_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, 4) };
559-
if r_inlined_data != l_inlined_data {
560-
return l_inlined_data.cmp(r_inlined_data);
559+
560+
// Note: In theory, ByteView is only used for string which is larger than 12 bytes,
561+
// but we can still use it to get the inlined prefix for shorter strings.
562+
// The prefix is always the first 4 bytes of the view, for both short and long strings.
563+
let l_inlined_be = l_byte_view.prefix.swap_bytes();
564+
let r_inlined_be = r_byte_view.prefix.swap_bytes();
565+
if l_inlined_be != r_inlined_be {
566+
return l_inlined_be.cmp(&r_inlined_be);
561567
}
562568

563569
// unfortunately, we need to compare the full data
@@ -566,6 +572,63 @@ impl<T: ByteViewType + ?Sized> GenericByteViewArray<T> {
566572

567573
l_full_data.cmp(r_full_data)
568574
}
575+
576+
/// Builds a 128-bit composite key for an inline value:
577+
///
578+
/// - High 96 bits: the inline data in big-endian byte order (for correct lexicographical sorting).
579+
/// - Low 32 bits: the length in big-endian byte order, acting as a tiebreaker so shorter strings
580+
/// (or those with fewer meaningful bytes) always numerically sort before longer ones.
581+
///
582+
/// This function extracts the length and the 12-byte inline string data from the raw
583+
/// little-endian `u128` representation, converts them to big-endian ordering, and packs them
584+
/// into a single `u128` value suitable for fast, branchless comparisons.
585+
///
586+
/// ### Why include length?
587+
///
588+
/// A pure 96-bit content comparison can’t distinguish between two values whose inline bytes
589+
/// compare equal—either because one is a true prefix of the other or because zero-padding
590+
/// hides extra bytes. By tucking the 32-bit length into the lower bits, a single `u128` compare
591+
/// handles both content and length in one go.
592+
///
593+
/// Example: comparing "bar" (3 bytes) vs "bar\0" (4 bytes)
594+
///
595+
/// | String | Bytes 0–4 (length LE) | Bytes 4–16 (data + padding) |
596+
/// |------------|-----------------------|---------------------------------|
597+
/// | `"bar"` | `03 00 00 00` | `62 61 72` + 9 × `00` |
598+
/// | `"bar\0"`| `04 00 00 00` | `62 61 72 00` + 8 × `00` |
599+
///
600+
/// Both inline parts become `62 61 72 00…00`, so they tie on content. The length field
601+
/// then differentiates:
602+
///
603+
/// ```text
604+
/// key("bar") = 0x0000000000000000000062617200000003
605+
/// key("bar\0") = 0x0000000000000000000062617200000004
606+
/// ⇒ key("bar") < key("bar\0")
607+
/// ```
608+
#[inline(always)]
609+
pub fn inline_key_fast(raw: u128) -> u128 {
610+
// Convert the raw u128 (little-endian) into bytes for manipulation
611+
let raw_bytes = raw.to_le_bytes();
612+
613+
// Extract the length (first 4 bytes), convert to big-endian u32, and promote to u128
614+
let len_le = &raw_bytes[0..4];
615+
let len_be = u32::from_le_bytes(len_le.try_into().unwrap()).to_be() as u128;
616+
617+
// Extract the inline string bytes (next 12 bytes), place them into the lower 12 bytes of a 16-byte array,
618+
// padding the upper 4 bytes with zero to form a little-endian u128 value
619+
let mut inline_bytes = [0u8; 16];
620+
inline_bytes[4..16].copy_from_slice(&raw_bytes[4..16]);
621+
622+
// Convert to big-endian to ensure correct lexical ordering
623+
let inline_u128 = u128::from_le_bytes(inline_bytes).to_be();
624+
625+
// Shift right by 32 bits to discard the zero padding (upper 4 bytes),
626+
// so that the inline string occupies the high 96 bits
627+
let inline_part = inline_u128 >> 32;
628+
629+
// Combine the inline string part (high 96 bits) and length (low 32 bits) into the final key
630+
(inline_part << 32) | len_be
631+
}
569632
}
570633

571634
impl<T: ByteViewType + ?Sized> Debug for GenericByteViewArray<T> {
@@ -874,7 +937,10 @@ impl From<Vec<Option<String>>> for StringViewArray {
874937
#[cfg(test)]
875938
mod tests {
876939
use crate::builder::{BinaryViewBuilder, StringViewBuilder};
877-
use crate::{Array, BinaryViewArray, StringViewArray};
940+
use crate::types::BinaryViewType;
941+
use crate::{
942+
Array, BinaryViewArray, GenericBinaryArray, GenericByteViewArray, StringViewArray,
943+
};
878944
use arrow_buffer::{Buffer, ScalarBuffer};
879945
use arrow_data::ByteView;
880946

@@ -1090,4 +1156,72 @@ mod tests {
10901156
assert_eq!(array2, array2.clone());
10911157
assert_eq!(array1, array2);
10921158
}
1159+
1160+
/// Integration tests for `inline_key_fast` covering:
1161+
///
1162+
/// 1. Monotonic ordering across increasing lengths and lexical variations.
1163+
/// 2. Cross-check against `GenericBinaryArray` comparison to ensure semantic equivalence.
1164+
///
1165+
/// This also includes a specific test for the “bar” vs. “bar\0” case, demonstrating why
1166+
/// the length field is required even when all inline bytes fit in 12 bytes.
1167+
#[test]
1168+
fn test_inline_key_fast_various_lengths_and_lexical() {
1169+
/// Helper to create a raw u128 value representing an inline ByteView
1170+
/// - `length`: number of meaningful bytes (≤ 12)
1171+
/// - `data`: the actual inline data
1172+
fn make_raw_inline(length: u32, data: &[u8]) -> u128 {
1173+
assert!(length as usize <= 12, "Inline length must be ≤ 12");
1174+
assert!(data.len() == length as usize, "Data must match length");
1175+
1176+
let mut raw_bytes = [0u8; 16];
1177+
raw_bytes[0..4].copy_from_slice(&length.to_le_bytes()); // little-endian length
1178+
raw_bytes[4..(4 + data.len())].copy_from_slice(data); // inline data
1179+
u128::from_le_bytes(raw_bytes)
1180+
}
1181+
1182+
// Test inputs: include the specific "bar" vs "bar\0" case, plus length and lexical variations
1183+
let test_inputs: Vec<&[u8]> = vec![
1184+
b"a",
1185+
b"aa",
1186+
b"aaa",
1187+
b"aab",
1188+
b"abcd",
1189+
b"abcde",
1190+
b"abcdef",
1191+
b"abcdefg",
1192+
b"abcdefgh",
1193+
b"abcdefghi",
1194+
b"abcdefghij",
1195+
b"abcdefghijk",
1196+
b"abcdefghijkl", // 12 bytes, max inline
1197+
b"bar",
1198+
b"bar\0", // special case to test length tiebreaker
1199+
b"xyy",
1200+
b"xyz",
1201+
];
1202+
1203+
// Monotonic key order: content then length,and cross-check against GenericBinaryArray comparison
1204+
let array: GenericBinaryArray<i32> =
1205+
GenericBinaryArray::from(test_inputs.iter().map(|s| Some(*s)).collect::<Vec<_>>());
1206+
1207+
for i in 0..array.len() - 1 {
1208+
let v1 = array.value(i);
1209+
let v2 = array.value(i + 1);
1210+
// Ensure lexical ordering matches
1211+
assert!(v1 < v2, "Array compare failed: {v1:?} !< {v2:?}");
1212+
// Ensure fast key compare matches
1213+
let key1 = GenericByteViewArray::<BinaryViewType>::inline_key_fast(make_raw_inline(
1214+
v1.len() as u32,
1215+
v1,
1216+
));
1217+
let key2 = GenericByteViewArray::<BinaryViewType>::inline_key_fast(make_raw_inline(
1218+
v2.len() as u32,
1219+
v2,
1220+
));
1221+
assert!(
1222+
key1 < key2,
1223+
"Key compare failed: key({v1:?})=0x{key1:032x} !< key({v2:?})=0x{key2:032x}",
1224+
);
1225+
}
1226+
}
10931227
}

arrow-ord/src/cmp.rs

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ use arrow_buffer::bit_util::ceil;
3333
use arrow_buffer::{BooleanBuffer, MutableBuffer, NullBuffer};
3434
use arrow_schema::ArrowError;
3535
use arrow_select::take::take;
36+
use std::cmp::Ordering;
3637
use std::ops::Not;
3738

3839
#[derive(Debug, Copy, Clone)]
@@ -571,7 +572,7 @@ impl<'a, T: ByteViewType> ArrayOrd for &'a GenericByteViewArray<T> {
571572
let r_view = unsafe { r.0.views().get_unchecked(r.1) };
572573
if l.0.data_buffers().is_empty() && r.0.data_buffers().is_empty() {
573574
// For eq case, we can directly compare the inlined bytes
574-
return l_view.cmp(r_view).is_eq();
575+
return l_view == r_view;
575576
}
576577

577578
let l_len = *l_view as u32;
@@ -592,15 +593,15 @@ impl<'a, T: ByteViewType> ArrayOrd for &'a GenericByteViewArray<T> {
592593

593594
#[inline(always)]
594595
fn is_lt(l: Self::Item, r: Self::Item) -> bool {
596+
// If both arrays use only the inline buffer
595597
if l.0.data_buffers().is_empty() && r.0.data_buffers().is_empty() {
596598
let l_view = unsafe { l.0.views().get_unchecked(l.1) };
597599
let r_view = unsafe { r.0.views().get_unchecked(r.1) };
598-
let l_len = *l_view as u32 as usize;
599-
let r_len = *r_view as u32 as usize;
600-
let l_bytes = unsafe { GenericByteViewArray::<T>::inline_value(l_view, l_len) };
601-
let r_bytes = unsafe { GenericByteViewArray::<T>::inline_value(r_view, r_len) };
602-
return l_bytes.cmp(r_bytes).is_lt();
600+
return GenericByteViewArray::<T>::inline_key_fast(*l_view)
601+
< GenericByteViewArray::<T>::inline_key_fast(*r_view);
603602
}
603+
604+
// Fallback to the generic, unchecked comparison for non-inline cases
604605
// # Safety
605606
// The index is within bounds as it is checked in value()
606607
unsafe { GenericByteViewArray::compare_unchecked(l.0, l.1, r.0, r.1).is_lt() }
@@ -642,17 +643,14 @@ pub fn compare_byte_view<T: ByteViewType>(
642643
left_idx: usize,
643644
right: &GenericByteViewArray<T>,
644645
right_idx: usize,
645-
) -> std::cmp::Ordering {
646+
) -> Ordering {
646647
assert!(left_idx < left.len());
647648
assert!(right_idx < right.len());
648649
if left.data_buffers().is_empty() && right.data_buffers().is_empty() {
649650
let l_view = unsafe { left.views().get_unchecked(left_idx) };
650651
let r_view = unsafe { right.views().get_unchecked(right_idx) };
651-
let l_len = *l_view as u32 as usize;
652-
let r_len = *r_view as u32 as usize;
653-
let l_bytes = unsafe { GenericByteViewArray::<T>::inline_value(l_view, l_len) };
654-
let r_bytes = unsafe { GenericByteViewArray::<T>::inline_value(r_view, r_len) };
655-
return l_bytes.cmp(r_bytes);
652+
return GenericByteViewArray::<T>::inline_key_fast(*l_view)
653+
.cmp(&GenericByteViewArray::<T>::inline_key_fast(*r_view));
656654
}
657655
unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) }
658656
}

0 commit comments

Comments
 (0)