This is an automated email from the ASF dual-hosted git repository.
alamb 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 803b51c38d perf: speed up substring_by_char with an ASCII fast path
and single-pass bounds (#10334)
803b51c38d is described below
commit 803b51c38d049b1355c89f592631289c49fbc629
Author: Andy Grove <[email protected]>
AuthorDate: Fri Jul 17 15:16:44 2026 -0600
perf: speed up substring_by_char with an ASCII fast path and single-pass
bounds (#10334)
# Which issue does this PR close?
N/A — performance improvement to an existing kernel.
# Rationale for this change
`substring_by_char` calls `val.chars().count()` on every value, which is
a full UTF-8 decode of the entire string. The result is only used when
`start < 0`, so for the common non-negative `start` case the work is
thrown away. When `start < 0`, the string is then walked a second time
from the front to convert the char index back into a byte offset.
It also decodes UTF-8 unconditionally, even when the array is entirely
ASCII, where a char index is a byte index and the bounds can be computed
arithmetically. The doc comment currently tells users to reach for
`substring` themselves in that case; the kernel can just detect it.
This is a port of the approach taken in apache/datafusion-comet#4903,
which replaced this kernel with a local copy for exactly these reasons.
# What changes are included in this PR?
- `substring_by_char` picks a bounds function once per array:
`ascii_bounds` (integer arithmetic on byte offsets) when `is_ascii()`
holds, `utf8_bounds` otherwise. `substring_by_char_impl` does the shared
buffer building.
- The unconditional `chars().count()` is gone. Negative starts use
`char_indices().nth_back()`, which decodes only as far back as needed
rather than scanning the whole value twice.
- `utf8_bounds` short-circuits the end scan when the requested char
length is at least the remaining byte length, since a char is never
smaller than one byte.
- The value buffer capacity is bounded by an upper estimate of the
output element size, so a short substring of long strings no longer
allocates the full input size.
- Semantics are unchanged, including the clamping of out-of-range starts
in both directions.
The output is still built with the checked `GenericStringArray::new`.
Skipping the redundant UTF-8 validation with `new_unchecked` is possible
(every slice is on a char boundary of an already-valid string) but is
left out of this PR.
Benchmark, 65,536 rows x 1,000-char strings:
| Benchmark | Before | After | Change |
|---|---|---|---|
| substring utf8 by char (existing) | 33.97 ms | 4.84 ms | -85.7% |
| ascii, prefix (start=0, length=10) | 2.61 ms | 2.34 ms | -10.1% |
| ascii, tail (start=-10) | 31.76 ms | 2.41 ms | -92.4% |
| non-ascii, prefix | 3.95 ms | 1.34 ms | -66.1% |
| non-ascii, tail | 65.62 ms | 1.06 ms | -98.4% |
The ASCII prefix case gains the least because it is dominated by the
copy and the output validation, but it still comes out ahead of the
`is_ascii()` scan it now pays for.
# Are these changes tested?
Yes. The existing `substring_by_char` tests all mix ASCII and non-ASCII
values in a single array, so they only ever exercised the UTF-8 path.
Added an ASCII-only test matrix (`ascii_string_by_char` /
`ascii_large_string_by_char`) covering identity, positive and negative
starts, out-of-range starts in both directions, zero length, and a
`u64::MAX` length to pin the saturating-add clamp.
Added four benchmarks to `substring_kernels.rs` (ASCII and non-ASCII,
prefix and tail) plus a non-ASCII array generator; the existing bench
name is unchanged so its history stays comparable.
# Are there any user-facing changes?
No API change. `substring_by_char` is faster, and its `# Performance`
doc note is updated to mention the ASCII fast path.
---
arrow-string/src/substring.rs | 153 ++++++++++++++++++++++++++++++-------
arrow/benches/substring_kernels.rs | 36 +++++++++
2 files changed, 160 insertions(+), 29 deletions(-)
diff --git a/arrow-string/src/substring.rs b/arrow-string/src/substring.rs
index 152f1f152d..c8762499ca 100644
--- a/arrow-string/src/substring.rs
+++ b/arrow-string/src/substring.rs
@@ -126,8 +126,9 @@ pub fn substring(
/// # Performance
///
/// This function is slower than [substring]. Theoretically, the time
complexity
-/// is `O(n)` where `n` is the length of the value buffer. It is recommended to
-/// use [substring] if the input array only contains ASCII chars.
+/// is `O(n)` where `n` is the length of the value buffer. If the array only
+/// contains ASCII chars, a fast path avoids decoding UTF-8 altogether and the
+/// performance is comparable to [substring].
///
/// # Basic usage
/// ```
@@ -142,57 +143,106 @@ pub fn substring_by_char<OffsetSize: OffsetSizeTrait>(
start: i64,
length: Option<u64>,
) -> Result<GenericStringArray<OffsetSize>, ArrowError> {
+ let length = length.map(|len| usize::try_from(len).unwrap_or(usize::MAX));
+
+ if array.is_ascii() {
+ // One char is one byte, so the byte offsets can be computed
arithmetically.
+ Ok(substring_by_char_impl(array, length, |val| {
+ ascii_bounds(val, start, length)
+ }))
+ } else {
+ // A char occupies at most 4 bytes.
+ let max_element_len = length.map(|len| len.saturating_mul(4));
+ Ok(substring_by_char_impl(array, max_element_len, |val| {
+ utf8_bounds(val, start, length)
+ }))
+ }
+}
+
+/// Builds the output of [`substring_by_char`], delegating to `bounds` to
locate the
+/// substring of each value.
+///
+/// * `max_element_len` - an upper bound, in bytes, on the size of a single
output
+/// element, used to avoid over-allocating the value buffer. [None] if
unbounded.
+fn substring_by_char_impl<OffsetSize: OffsetSizeTrait, F: Fn(&str) -> (usize,
usize)>(
+ array: &GenericStringArray<OffsetSize>,
+ max_element_len: Option<usize>,
+ bounds: F,
+) -> GenericStringArray<OffsetSize> {
let mut vals = BufferBuilder::<u8>::new({
let offsets = array.value_offsets();
- (offsets[array.len()] - offsets[0]).to_usize().unwrap()
+ let input_len = (offsets[array.len()] -
offsets[0]).to_usize().unwrap();
+ match max_element_len {
+ Some(len) => input_len.min(array.len().saturating_mul(len)),
+ None => input_len,
+ }
});
let mut new_offsets = BufferBuilder::<OffsetSize>::new(array.len() + 1);
new_offsets.append(OffsetSize::zero());
- let length = length.map(|len| len.to_usize().unwrap());
array.iter().for_each(|val| {
if let Some(val) = val {
- let char_count = val.chars().count();
- let start = if start >= 0 {
- start.to_usize().unwrap()
- } else {
- char_count - (-start).to_usize().unwrap().min(char_count)
- };
- let (start_offset, end_offset) = get_start_end_offset(val, start,
length);
+ let (start_offset, end_offset) = bounds(val);
vals.append_slice(&val.as_bytes()[start_offset..end_offset]);
}
new_offsets.append(OffsetSize::from_usize(vals.len()).unwrap());
});
+
let offsets = OffsetBuffer::new(new_offsets.finish().into());
let values = vals.finish();
let nulls = array
.nulls()
.map(|n| n.inner().sliced())
.and_then(|b| NullBuffer::from_unsliced_buffer(b, array.len()));
- Ok(GenericStringArray::<OffsetSize>::new(
- offsets, values, nulls,
- ))
+ GenericStringArray::<OffsetSize>::new(offsets, values, nulls)
}
-/// * `val` - string
-/// * `start` - the start char index of the substring
-/// * `length` - the char length of the substring
+/// Returns the `start` and `end` byte offset of the substring of an ASCII
`val`, where
+/// one char is one byte.
///
-/// Return the `start` and `end` offset (by byte) of the substring
-fn get_start_end_offset(val: &str, start: usize, length: Option<usize>) ->
(usize, usize) {
+/// * `start` - the start char index of the substring, counted from the end if
negative
+/// * `length` - the char length of the substring, or [None] to take the rest
of `val`
+#[inline]
+fn ascii_bounds(val: &str, start: i64, length: Option<usize>) -> (usize,
usize) {
let len = val.len();
- let mut offset_char_iter = val.char_indices();
- let start_offset = offset_char_iter
- .nth(start)
- .map_or(len, |(offset, _)| offset);
+ let start_offset = if start >= 0 {
+ usize::try_from(start).unwrap_or(usize::MAX).min(len)
+ } else {
+
len.saturating_sub(usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX))
+ };
+ let end_offset = length.map_or(len, |length|
start_offset.saturating_add(length).min(len));
+ (start_offset, end_offset)
+}
+
+/// Returns the `start` and `end` byte offset of the substring of an arbitrary
UTF-8 `val`.
+///
+/// * `start` - the start char index of the substring, counted from the end if
negative
+/// * `length` - the char length of the substring, or [None] to take the rest
of `val`
+#[inline]
+fn utf8_bounds(val: &str, start: i64, length: Option<usize>) -> (usize, usize)
{
+ let len = val.len();
+ let start_offset = if start >= 0 {
+ val.char_indices()
+ .nth(usize::try_from(start).unwrap_or(usize::MAX))
+ .map_or(len, |(offset, _)| offset)
+ } else {
+ // `start` is negative, so `nth_back` counts back from the last char.
Strings with
+ // fewer than `-start` chars start at 0.
+ let back = usize::try_from(start.unsigned_abs()).unwrap_or(usize::MAX);
+ val.char_indices()
+ .nth_back(back - 1)
+ .map_or(0, |(offset, _)| offset)
+ };
let end_offset = length.map_or(len, |length| {
- if length > 0 {
- offset_char_iter
- .nth(length - 1)
- .map_or(len, |(offset, _)| offset)
- } else {
- start_offset
+ // Every char is at least one byte, so the rest of `val` is shorter
than `length`
+ // chars and there is nothing to scan for.
+ if length >= len - start_offset {
+ return len;
}
+ val[start_offset..]
+ .char_indices()
+ .nth(length)
+ .map_or(len, |(offset, _)| start_offset + offset)
});
(start_offset, end_offset)
}
@@ -828,6 +878,51 @@ mod tests {
without_nulls_generic_string_by_char::<i64>()
}
+ /// The array is all-ASCII, which takes the fast path of
[`substring_by_char`].
+ fn ascii_generic_string_by_char<O: OffsetSizeTrait>() {
+ let input = vec![Some("hello"), None, Some(""), Some("rust")];
+ let cases = gen_test_cases!(
+ input,
+ // identity
+ (0, None, input.clone()),
+ // increase start
+ (1, None, vec![Some("ello"), None, Some(""), Some("ust")]),
+ (4, None, vec![Some("o"), None, Some(""), Some("")]),
+ // high start -> nothing
+ (1000, None, vec![Some(""), None, Some(""), Some("")]),
+ // increase start negatively
+ (-1, None, vec![Some("o"), None, Some(""), Some("t")]),
+ (-4, None, vec![Some("ello"), None, Some(""), Some("rust")]),
+ // high negative start -> identity
+ (-1000, None, input.clone()),
+ // 0 length -> nothing
+ (0, Some(0), vec![Some(""), None, Some(""), Some("")]),
+ // increase length
+ (1, Some(2), vec![Some("el"), None, Some(""), Some("us")]),
+ (-4, Some(2), vec![Some("el"), None, Some(""), Some("ru")]),
+ // high length -> identity
+ (0, Some(1000), input.clone()),
+ // length past the end is clamped, including when it would overflow
+ (
+ 1,
+ Some(u64::MAX),
+ vec![Some("ello"), None, Some(""), Some("ust")]
+ )
+ );
+
+ do_test!(cases, GenericStringArray<O>, substring_by_char);
+ }
+
+ #[test]
+ fn ascii_string_by_char() {
+ ascii_generic_string_by_char::<i32>()
+ }
+
+ #[test]
+ fn ascii_large_string_by_char() {
+ ascii_generic_string_by_char::<i64>()
+ }
+
fn generic_string_by_char_with_non_zero_offset<O: OffsetSizeTrait>() {
let values = "S→T = Πx:S.T";
let offsets = &[
diff --git a/arrow/benches/substring_kernels.rs
b/arrow/benches/substring_kernels.rs
index bf910e7aaf..a3b58988c4 100644
--- a/arrow/benches/substring_kernels.rs
+++ b/arrow/benches/substring_kernels.rs
@@ -38,11 +38,31 @@ fn bench_substring_by_char<O: OffsetSizeTrait>(
substring_by_char(hint::black_box(arr), start, length).unwrap();
}
+/// A non-ASCII variant of [`create_string_array_with_len`]: every fourth char
is a
+/// multi-byte codepoint, so `substring_by_char` has to decode UTF-8 to find a
char index.
+fn create_utf8_array_with_len<O: OffsetSizeTrait>(
+ size: usize,
+ str_len: usize,
+) -> GenericStringArray<O> {
+ create_string_array_with_len::<O>(size, 0.0, str_len)
+ .iter()
+ .map(|val| {
+ val.map(|val| {
+ val.chars()
+ .enumerate()
+ .map(|(i, c)| if i % 4 == 0 { 'é' } else { c })
+ .collect::<String>()
+ })
+ })
+ .collect()
+}
+
fn add_benchmark(c: &mut Criterion) {
let size = 65536;
let val_len = 1000;
let arr_string = create_string_array_with_len::<i32>(size, 0.0, val_len);
+ let arr_utf8 = create_utf8_array_with_len::<i32>(size, val_len);
let arr_fsb = create_fsb_array(size, 0.0, val_len);
c.bench_function("substring utf8 (start = 0, length = None)", |b| {
@@ -57,6 +77,22 @@ fn add_benchmark(c: &mut Criterion) {
b.iter(|| bench_substring_by_char(&arr_string, 1, Some((val_len - 1)
as u64)))
});
+ c.bench_function("substring by char (ascii, prefix)", |b| {
+ b.iter(|| bench_substring_by_char(&arr_string, 0, Some(10)))
+ });
+
+ c.bench_function("substring by char (ascii, tail)", |b| {
+ b.iter(|| bench_substring_by_char(&arr_string, -10, None))
+ });
+
+ c.bench_function("substring by char (non-ascii, prefix)", |b| {
+ b.iter(|| bench_substring_by_char(&arr_utf8, 0, Some(10)))
+ });
+
+ c.bench_function("substring by char (non-ascii, tail)", |b| {
+ b.iter(|| bench_substring_by_char(&arr_utf8, -10, None))
+ });
+
c.bench_function("substring fixed size binary array", |b| {
b.iter(|| bench_substring(&arr_fsb, 1, Some((val_len - 1) as u64)))
});
