perf: Optimize left, right

datafusion/functions/src/strings.rs

@@ -17,6 +17,7 @@
 
 use std::marker::PhantomData;
 use std::mem::size_of;
+use std::ops::Range;
 use std::sync::Arc;
 
 use datafusion_common::utils::offset_span_len;
@@ -27,8 +28,9 @@ use datafusion_common::{
 
 use arrow::array::{
     Array, ArrayAccessor, ArrayDataBuilder, ArrayRef, BinaryArray, BinaryViewArray,
-    ByteView, GenericStringArray, LargeBinaryArray, LargeStringArray, OffsetSizeTrait,
-    StringArray, StringViewArray, as_largestring_array, make_view,
+    ByteView, GenericStringArray, LargeBinaryArray, LargeStringArray,
+    MAX_INLINE_VIEW_LEN, OffsetSizeTrait, StringArray, StringViewArray,
+    as_largestring_array, make_view,
 };
 use arrow::buffer::{Buffer, MutableBuffer, NullBuffer, ScalarBuffer};
 use arrow::datatypes::DataType;
@@ -1217,6 +1219,73 @@ pub(crate) fn append_view(
     views_buffer.push(sub_view);
 }
 
+/// Values of at most this many bytes are stored inline in their view.
+pub(crate) const MAX_INLINE_LEN: usize = MAX_INLINE_VIEW_LEN as usize;
+
+/// Returns the view for a substring of the value an existing view refers to.
+///
+/// # Arguments
+/// - view: The original view value
+/// - source: The bytes the original view refers to
+/// - range: The byte range within `source` of the substring to return a view for
+///
+/// Substrings longer than 12 bytes point into the same data buffer as `view`;
+/// shorter ones are stored inline in the new view.
+///
+/// This uses shifts and masks rather than [`make_view`], which picks copy code
+/// based on the substring's length; the CPU often mispredicts that choice when
+/// lengths vary from row to row.
+#[inline]
+pub(crate) fn sub_view(view: u128, source: &[u8], range: Range<usize>) -> u128 {
+    debug_assert!(range.start <= range.end && range.end <= source.len());
+
+    // The substring's first bytes, in the low-order bits. Any bits past the
+    // end of the substring are masked off by `inline_view`.
+    let leading_bytes = if source.len() <= MAX_INLINE_LEN {
+        // `source` is stored in `view` itself, after its 4-byte length.
+        (view >> 32) >> (8 * range.start)
+    } else {
+        // `source` has more than 12 bytes, so read the 12 bytes starting at
+        // `range.start`, or the last 12 bytes if that would run past the end,
+        // and skip any that come before `range.start`.
+        let window_start = range.start.min(source.len() - MAX_INLINE_LEN);
+        let window = source[window_start..window_start + MAX_INLINE_LEN]
+            .try_into()
+            .unwrap();
+        read_12_bytes(window) >> (8 * (range.start - window_start))
+    };
+
+    let len = range.len();
+    if len <= MAX_INLINE_LEN {
+        inline_view(leading_bytes, len)
+    } else {
+        let original = ByteView::from(view);
+        ByteView {
+            length: len as u32,
+            prefix: leading_bytes as u32,
+            offset: original.offset + range.start as u32,
+            ..original
+        }
+        .as_u128()
+    }
+}
+
+/// Reads `bytes` as a little-endian integer.
+#[inline]
+fn read_12_bytes(bytes: &[u8; 12]) -> u128 {
+    let low = u64::from_le_bytes(bytes[..8].try_into().unwrap());
+    let high = u32::from_le_bytes(bytes[8..].try_into().unwrap());
+    u128::from(low) | (u128::from(high) << 64)
+}
+
+/// Returns an inline view of the `len` low-order bytes of `bytes`.
+#[inline]
+fn inline_view(bytes: u128, len: usize) -> u128 {
+    debug_assert!(len <= MAX_INLINE_LEN);
+    let mask = (1u128 << (8 * len)) - 1;
+    ((bytes & mask) << 32) | len as u128
+}
+
 #[derive(Debug)]
 pub(crate) enum ColumnarValueRef<'a> {
     Scalar(&'a [u8]),
@@ -1430,6 +1499,31 @@ mod tests {
     use super::*;
     use crate::utils::test::sliced_byte_array;
 
+    #[test]
+    fn sub_view_matches_make_view() {
+        let buffer_index = 3;
+        let offset = 100;
+        for source_len in 0..=40 {
+            // Include bytes >= 0x80 so that masking errors would show up.
+            let source: Vec<u8> = (0..source_len).map(|i| (i * 37 + 11) as u8).collect();
+            let view = make_view(&source, buffer_index, offset);
+            for start in 0..=source_len {
+                for end in start..=source_len {
+                    let expected = make_view(
+                        &source[start..end],
+                        buffer_index,
+                        offset + start as u32,
+                    );
+                    assert_eq!(
+                        sub_view(view, &source, start..end),
+                        expected,
+                        "source_len={source_len}, range={start}..{end}"
+                    );
+                }
+            }
+        }
+    }
+
     #[test]
     fn concat_capacity_uses_visible_bytes() -> Result<()> {
         for data_type in [

datafusion/functions/src/unicode/common.rs

@@ -17,9 +17,9 @@
 
 //! Common utilities for implementing unicode functions
 
+use crate::strings::{MAX_INLINE_LEN, sub_view};
 use arrow::array::{
-    Array, ArrayRef, ByteView, GenericStringArray, Int64Array, OffsetSizeTrait,
-    StringViewArray, make_view,
+    Array, ArrayRef, GenericStringArray, Int64Array, OffsetSizeTrait, StringViewArray,
 };
 use arrow::datatypes::DataType;
 use arrow_buffer::{NullBuffer, ScalarBuffer};
@@ -126,21 +126,34 @@ pub(crate) enum StringCharLen {
 #[inline]
 fn left_right_byte_length(string: &str, n: i64) -> usize {
     let abs = n.unsigned_abs().min(usize::MAX as u64) as usize;
-    // For ASCII input every character is exactly one byte, so the byte offset of
-    // the n-th codepoint is just the (clamped) character count. This avoids the
-    // per-character `char_indices()` scan of the general path.
+    let bytes = string.as_bytes();
+    // ASCII bytes are never part of a multi-byte UTF-8 sequence, so if the
+    // `abs` bytes at the relevant end of the string are ASCII, they are exactly
+    // the `abs` characters at that end. Checking only those bytes is cheaper
+    // than either a `char_indices()` scan or checking the whole string.
     match n.cmp(&0) {
         Ordering::Equal => 0,
-        // `abs` chars trimmed from the end: keep the leading `len - abs`.
-        Ordering::Less if string.is_ascii() => string.len().saturating_sub(abs),
-        Ordering::Less => string
-            .char_indices()
-            .nth_back(abs - 1)
-            .map(|(index, _)| index)
-            .unwrap_or(0),
-        // First `abs` chars, but never past the end of the string.
-        Ordering::Greater if string.is_ascii() => abs.min(string.len()),
-        Ordering::Greater => byte_offset_of_char(string, abs),
+        // Byte offset of the `abs`-th character from the end.
+        Ordering::Less => {
+            let start = bytes.len().saturating_sub(abs);
+            if bytes[start..].is_ascii() {
+                start
+            } else {
+                string
+                    .char_indices()
+                    .nth_back(abs - 1)
+                    .map_or(0, |(index, _)| index)
+            }
+        }
+        // Byte offset of the `abs`-th character from the start.
+        Ordering::Greater => {
+            let end = abs.min(bytes.len());
+            if bytes[..end].is_ascii() {
+                end
+            } else {
+                byte_offset_of_char(string, abs)
+            }
+        }
     }
 }
 
@@ -204,14 +217,10 @@ fn general_left_right_view<F: LeftRightSlicer>(
             let n = n_array.value(idx);
 
             let range = F::slice(string, n);
-            let result_bytes = &string.as_bytes()[range.clone()];
-            if result_bytes.len() > 12 {
+            if range.len() > MAX_INLINE_LEN {
                 has_out_of_line = true;
             }
-
-            let byte_view = ByteView::from(views[idx]);
-            let new_offset = byte_view.offset + (range.start as u32);
-            make_view(result_bytes, byte_view.buffer_index, new_offset)
+            sub_view(views[idx], string.as_bytes(), range)
         })
         .collect::<Vec<u128>>();
 
@@ -223,7 +232,10 @@ fn general_left_right_view<F: LeftRightSlicer>(
     };
 
     // SAFETY:
-    // - Each view is produced by `make_view` with correct bytes and offset
+    // - Each view is produced by `sub_view` from the input view and a range
+    //   within the input string, as returned by `F::slice`
+    // - `F::slice` returns ranges that start and end on char boundaries (see
+    //   `left_right_byte_length`), so every result is valid UTF-8
     // - Out-of-line views reuse the original buffer index and adjusted offset
     unsafe {
         let array = StringViewArray::new_unchecked(views, data_buffers, new_nulls);