diff --git a/datafusion/functions/src/strings.rs b/datafusion/functions/src/strings.rs index 9bdb858ba1b62..96084474fff9e 100644 --- a/datafusion/functions/src/strings.rs +++ b/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) -> 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 = (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 [ diff --git a/datafusion/functions/src/unicode/common.rs b/datafusion/functions/src/unicode/common.rs index ac24877863e80..a9074b9fdf606 100644 --- a/datafusion/functions/src/unicode/common.rs +++ b/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( 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::>(); @@ -223,7 +232,10 @@ fn general_left_right_view( }; // 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);