From f1f45133baa67a90781305f9cdb3a1eed12df4a6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 13 Jul 2026 16:24:37 -0600 Subject: [PATCH] perf: optimize regexp_instr in datafusion-functions --- datafusion/functions/Cargo.toml | 5 + datafusion/functions/benches/regexp_instr.rs | 99 ++++++++ datafusion/functions/src/regex/regexpinstr.rs | 235 +++++++++--------- 3 files changed, 215 insertions(+), 124 deletions(-) create mode 100644 datafusion/functions/benches/regexp_instr.rs diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 94830ee360585..dc502352c34a6 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -326,6 +326,11 @@ harness = false name = "regexp_count" required-features = ["regex_expressions"] +[[bench]] +harness = false +name = "regexp_instr" +required-features = ["regex_expressions"] + [[bench]] harness = false name = "crypto" diff --git a/datafusion/functions/benches/regexp_instr.rs b/datafusion/functions/benches/regexp_instr.rs new file mode 100644 index 0000000000000..9ac630d8c4b6e --- /dev/null +++ b/datafusion/functions/benches/regexp_instr.rs @@ -0,0 +1,99 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::Int64Array; +use arrow::array::OffsetSizeTrait; +use arrow::datatypes::{DataType, Field}; +use arrow::util::bench_util::create_string_array_with_len; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{DataFusionError, ScalarValue}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions::regex; +use std::hint::black_box; +use std::sync::Arc; + +fn create_args( + size: usize, + str_len: usize, + with_start: bool, +) -> Vec { + let string_array = Arc::new(create_string_array_with_len::(size, 0.1, str_len)); + let pattern = ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))); + + if with_start { + let start_array = Arc::new(Int64Array::from( + (0..size).map(|i| (i % 10 + 1) as i64).collect::>(), + )); + vec![ + ColumnarValue::Array(string_array), + pattern, + ColumnarValue::Array(start_array), + ] + } else { + vec![ColumnarValue::Array(string_array), pattern] + } +} + +fn invoke_regexp_instr_with_args( + args: Vec, + number_rows: usize, +) -> Result { + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) + .collect::>(); + let config_options = Arc::new(ConfigOptions::default()); + + regex::regexp_instr().invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Field::new("f", DataType::Int64, true).into(), + config_options: Arc::clone(&config_options), + }) +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 1024; + + for str_len in [32, 128] { + let args = create_args::(size, str_len, false); + c.bench_function( + &format!("regexp_instr_no_start [size={size}, str_len={str_len}]"), + |b| { + b.iter(|| { + black_box(invoke_regexp_instr_with_args(args.clone(), size).unwrap()) + }) + }, + ); + + let args = create_args::(size, str_len, true); + c.bench_function( + &format!("regexp_instr_with_start [size={size}, str_len={str_len}]"), + |b| { + b.iter(|| { + black_box(invoke_regexp_instr_with_args(args.clone(), size).unwrap()) + }) + }, + ); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/regex/regexpinstr.rs b/datafusion/functions/src/regex/regexpinstr.rs index d46e4452dbab1..9d62fe2ffb3c2 100644 --- a/datafusion/functions/src/regex/regexpinstr.rs +++ b/datafusion/functions/src/regex/regexpinstr.rs @@ -16,7 +16,7 @@ // under the License. use arrow::array::{ - Array, ArrayRef, AsArray, Datum, Int64Array, PrimitiveArray, StringArrayType, + Array, ArrayRef, AsArray, Datum, Int64Array, Int64Builder, StringArrayType, }; use arrow::datatypes::{DataType, Int64Type}; use arrow::datatypes::{ @@ -29,12 +29,12 @@ use datafusion_expr::{ TypeSignature::Exact, TypeSignature::Uniform, Volatility, }; use datafusion_macros::user_doc; -use itertools::izip; use regex::Regex; use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::sync::Arc; -use crate::regex::compile_and_cache_regex; +use crate::regex::compile_regex; #[user_doc( doc_section(label = "Regular Expression Functions"), @@ -240,7 +240,7 @@ fn regexp_instr( ®ex_array.as_string::(), start_array.map(|start| start.as_primitive::()), nth_array.map(|nth| nth.as_primitive::()), - Some(flags_array.as_string::()), + Some(&flags_array.as_string::()), subexpr_array.map(|subexpr| subexpr.as_primitive::()), ), (LargeUtf8, LargeUtf8, None) => regexp_instr_inner( @@ -256,7 +256,7 @@ fn regexp_instr( ®ex_array.as_string::(), start_array.map(|start| start.as_primitive::()), nth_array.map(|nth| nth.as_primitive::()), - Some(flags_array.as_string::()), + Some(&flags_array.as_string::()), subexpr_array.map(|subexpr| subexpr.as_primitive::()), ), (Utf8View, Utf8View, None) => regexp_instr_inner( @@ -272,7 +272,7 @@ fn regexp_instr( ®ex_array.as_string_view(), start_array.map(|start| start.as_primitive::()), nth_array.map(|nth| nth.as_primitive::()), - Some(flags_array.as_string_view()), + Some(&flags_array.as_string_view()), subexpr_array.map(|subexpr| subexpr.as_primitive::()), ), _ => Err(ArrowError::ComputeError( @@ -286,120 +286,104 @@ fn regexp_instr_inner<'a, S>( regex_array: &S, start_array: Option<&Int64Array>, nth_array: Option<&Int64Array>, - flags_array: Option, + flags_array: Option<&S>, subexp_array: Option<&Int64Array>, ) -> Result where S: StringArrayType<'a>, { let len = values.len(); + let mut regex_cache = RegexCache::default(); + let mut result = Int64Builder::with_capacity(len); - let default_start_array = PrimitiveArray::::from(vec![1; len]); - let start_array = start_array.unwrap_or(&default_start_array); - let start_input: Vec = (0..start_array.len()) - .map(|i| start_array.value(i)) // handle nulls as 0 - .collect(); - - let default_nth_array = PrimitiveArray::::from(vec![1; len]); - let nth_array = nth_array.unwrap_or(&default_nth_array); - let nth_input: Vec = (0..nth_array.len()) - .map(|i| nth_array.value(i)) // handle nulls as 0 - .collect(); - - let flags_input = match flags_array { - Some(flags) => flags.iter().collect(), - None => vec![None; len], - }; + for i in 0..len { + if regex_array.is_null(i) { + result.append_null(); + continue; + } + let regex = regex_array.value(i); + if regex.is_empty() { + result.append_value(0); + continue; + } - let default_subexp_array = PrimitiveArray::::from(vec![0; len]); - let subexp_array = subexp_array.unwrap_or(&default_subexp_array); - let subexp_input: Vec = (0..subexp_array.len()) - .map(|i| subexp_array.value(i)) // handle nulls as 0 - .collect(); - - let mut regex_cache = HashMap::new(); - - let result: Result>, ArrowError> = izip!( - values.iter(), - regex_array.iter(), - start_input.iter(), - nth_input.iter(), - flags_input.iter(), - subexp_input.iter() - ) - .map(|(value, regex, start, nth, flags, subexp)| match regex { - None => Ok(None), - Some("") => Ok(Some(0)), - Some(regex) => get_index( - value, - regex, - *start, - *nth, - *subexp, - *flags, - &mut regex_cache, - ), - }) - .collect(); - Ok(Arc::new(Int64Array::from(result?))) -} + if values.is_null(i) { + result.append_null(); + continue; + } + let value = values.value(i); + if value.is_empty() { + result.append_value(0); + continue; + } -fn handle_subexp( - pattern: &Regex, - search_slice: &str, - subexpr: i64, - value: &str, - byte_start_offset: usize, -) -> Result, ArrowError> { - if let Some(captures) = pattern.captures(search_slice) - && let Some(matched) = captures.get(subexpr as usize) - { - // Convert byte offset relative to search_slice back to 1-based character offset - // relative to the original `value` string. - let start_char_offset = - value[..byte_start_offset + matched.start()].chars().count() as i64 + 1; - return Ok(Some(start_char_offset)); + let flags = match flags_array { + Some(flags) if !flags.is_null(i) => Some(flags.value(i)), + _ => None, + }; + let pattern = regex_cache.get_or_compile(regex, flags)?; + + // The defaults apply when the optional argument was not supplied at + // all. A supplied but null slot reads through as its raw buffer value. + let start = start_array.map_or(1, |array| array.value(i)); + let nth = nth_array.map_or(1, |array| array.value(i)); + let subexp = subexp_array.map_or(0, |array| array.value(i)); + + result.append_value(get_index(value, pattern, start, nth, subexp)?); } - Ok(Some(0)) // Return 0 if the subexpression was not found + + Ok(Arc::new(result.finish())) } -fn get_nth_match( - pattern: &Regex, - search_slice: &str, - n: i64, - byte_start_offset: usize, - value: &str, -) -> Result, ArrowError> { - if let Some(mat) = pattern.find_iter(search_slice).nth((n - 1) as usize) { - // Convert byte offset relative to search_slice back to 1-based character offset - // relative to the original `value` string. - let match_start_byte_offset = byte_start_offset + mat.start(); - let match_start_char_offset = - value[..match_start_byte_offset].chars().count() as i64 + 1; - Ok(Some(match_start_char_offset)) - } else { - Ok(Some(0)) // Return 0 if the N-th match was not found +/// Compiles the patterns seen so far, keyed by `(pattern, flags)`. +/// +/// Patterns are addressed by index rather than by reference so that `last` can +/// memoize the previous row's pattern without holding a borrow of `indices` +/// across rows. A literal pattern yields the same string on every row, so that +/// memo means the common case never hashes a key. +#[derive(Default)] +struct RegexCache<'a> { + compiled: Vec, + indices: HashMap<(&'a str, Option<&'a str>), usize>, + last: Option<((&'a str, Option<&'a str>), usize)>, +} + +impl<'a> RegexCache<'a> { + fn get_or_compile( + &mut self, + regex: &'a str, + flags: Option<&'a str>, + ) -> Result<&Regex, ArrowError> { + let key = (regex, flags); + let index = match self.last { + Some((last_key, index)) if last_key == key => index, + _ => { + let index = match self.indices.entry(key) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + self.compiled.push(compile_regex(regex, flags)?); + *entry.insert(self.compiled.len() - 1) + } + }; + self.last = Some((key, index)); + index + } + }; + Ok(&self.compiled[index]) } } -fn get_index<'strings, 'cache>( - value: Option<&str>, - pattern: &'strings str, + +/// Returns the 1-based character position of the `n`-th match of `pattern` in +/// `value`, or 0 if there is no such match. The search begins at the 1-based +/// character position `start`. A positive `subexpr` selects that capture group +/// of the first match instead of the `n`-th match. `value` is non-empty. +fn get_index( + value: &str, + pattern: &Regex, start: i64, n: i64, subexpr: i64, - flags: Option<&'strings str>, - regex_cache: &'cache mut HashMap<(&'strings str, Option<&'strings str>), Regex>, -) -> Result, ArrowError> -where - 'strings: 'cache, -{ - let value = match value { - None => return Ok(None), - Some("") => return Ok(Some(0)), - Some(value) => value, - }; - let pattern: &Regex = compile_and_cache_regex(pattern, flags, regex_cache)?; - // println!("get_index: value = {}, pattern = {}, start = {}, n = {}, subexpr = {}, flags = {:?}", value, pattern, start, n, subexpr, flags); +) -> Result { if start < 1 { return Err(ArrowError::ComputeError( "regexp_instr() requires start to be 1-based".to_string(), @@ -412,31 +396,33 @@ where )); } - // --- Simplified byte_start_offset calculation --- - let total_chars = value.chars().count() as i64; - let byte_start_offset: usize = if start > total_chars { - // If start is beyond the total characters, it means we start searching - // after the string effectively. No matches possible. - return Ok(Some(0)); - } else { - // Get the byte offset for the (start - 1)-th character (0-based) - value - .char_indices() - .nth((start - 1) as usize) - .map(|(idx, _)| idx) - .unwrap_or(0) // Should not happen if start is valid and <= total_chars + // Byte offset of the `start`-th character. A `start` past the end of the + // string leaves nothing to search, so no match is possible. + let Some((byte_start_offset, _)) = value.char_indices().nth((start - 1) as usize) + else { + return Ok(0); }; - // --- End simplified calculation --- - let search_slice = &value[byte_start_offset..]; - // Handle subexpression capturing first, as it takes precedence - if subexpr > 0 { - return handle_subexp(pattern, search_slice, subexpr, value, byte_start_offset); - } + // A subexpression, when requested, takes precedence over the N-th match. + let match_start = if subexpr > 0 { + pattern + .captures(search_slice) + .and_then(|captures| captures.get(subexpr as usize)) + .map(|matched| matched.start()) + } else { + // `n` is 1-based, `nth` is 0-based. + pattern + .find_iter(search_slice) + .nth((n - 1) as usize) + .map(|matched| matched.start()) + }; - // Use nth to get the N-th match (n is 1-based, nth is 0-based) - get_nth_match(pattern, search_slice, n, byte_start_offset, value) + // Convert the byte offset within `search_slice` back to a 1-based character + // offset within `value`. + Ok(match_start.map_or(0, |offset| { + value[..byte_start_offset + offset].chars().count() as i64 + 1 + })) } #[cfg(test)] @@ -445,6 +431,7 @@ mod tests { use arrow::array::{GenericStringArray, StringViewArray}; use arrow::datatypes::Field; use datafusion_common::config::ConfigOptions; + use itertools::izip; #[test] fn test_regexp_instr() { test_case_sensitive_regexp_instr_nulls();