diff --git a/arrow-string/src/substring.rs b/arrow-string/src/substring.rs index 152f1f152dd4..c8762499ca78 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( start: i64, length: Option, ) -> Result, 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 (usize, usize)>( + array: &GenericStringArray, + max_element_len: Option, + bounds: F, +) -> GenericStringArray { let mut vals = BufferBuilder::::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::::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::::new( - offsets, values, nulls, - )) + GenericStringArray::::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) { +/// * `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) { 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) { + 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::() } + /// The array is all-ASCII, which takes the fast path of [`substring_by_char`]. + fn ascii_generic_string_by_char() { + 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, substring_by_char); + } + + #[test] + fn ascii_string_by_char() { + ascii_generic_string_by_char::() + } + + #[test] + fn ascii_large_string_by_char() { + ascii_generic_string_by_char::() + } + fn generic_string_by_char_with_non_zero_offset() { let values = "S→T = Πx:S.T"; let offsets = &[ diff --git a/arrow/benches/substring_kernels.rs b/arrow/benches/substring_kernels.rs index bf910e7aaf0a..a3b58988c408 100644 --- a/arrow/benches/substring_kernels.rs +++ b/arrow/benches/substring_kernels.rs @@ -38,11 +38,31 @@ fn bench_substring_by_char( 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( + size: usize, + str_len: usize, +) -> GenericStringArray { + create_string_array_with_len::(size, 0.0, str_len) + .iter() + .map(|val| { + val.map(|val| { + val.chars() + .enumerate() + .map(|(i, c)| if i % 4 == 0 { 'é' } else { c }) + .collect::() + }) + }) + .collect() +} + fn add_benchmark(c: &mut Criterion) { let size = 65536; let val_len = 1000; let arr_string = create_string_array_with_len::(size, 0.0, val_len); + let arr_utf8 = create_utf8_array_with_len::(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))) });