perf: speed up substring_by_char with an ASCII fast path and single-pass bounds

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 = &[

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)))
     });