From b4d9152367c8a0233beaa8b9817a69c0969e0c11 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 1 Sep 2026 21:59:08 +0700 Subject: [PATCH 1/3] perf: compile user regex patterns once per planned expression regexp_extract, regexp_extract_all, and split compiled the user pattern with Regex::new inside the per-batch evaluation path, so every 8192-row batch paid a full regex compile. The pattern cannot be hoisted to construction time because these are scalar functions created by name, with the pattern arriving per invocation as a scalar argument. Each planned expression now owns a one-slot pattern cache that compiles only when the pattern string changes, the same cost model rlike already has. Error messages and the phase at which an invalid pattern fails are unchanged. regexp_extract drops from 862us to 705us per 8192-row batch on the criterion bench, and a small-batch run (512 rows, 5000 batches) is 2.1x faster. split is unchanged on literal delimiters, which never compile a regex. --- native/spark-expr/benches/regexp_extract.rs | 11 +- .../spark-expr/benches/regexp_extract_all.rs | 8 +- native/spark-expr/benches/split.rs | 8 +- native/spark-expr/src/comet_scalar_funcs.rs | 14 ++- native/spark-expr/src/string_funcs/mod.rs | 2 + .../src/string_funcs/pattern_cache.rs | 106 ++++++++++++++++++ .../src/string_funcs/regexp_extract.rs | 46 ++++++-- .../src/string_funcs/regexp_extract_all.rs | 53 +++++++-- .../src/string_funcs/regexp_extract_common.rs | 6 +- native/spark-expr/src/string_funcs/split.rs | 86 +++++++++++--- 10 files changed, 296 insertions(+), 44 deletions(-) create mode 100644 native/spark-expr/src/string_funcs/pattern_cache.rs diff --git a/native/spark-expr/benches/regexp_extract.rs b/native/spark-expr/benches/regexp_extract.rs index 31cf9f32a74..9f514116c96 100644 --- a/native/spark-expr/benches/regexp_extract.rs +++ b/native/spark-expr/benches/regexp_extract.rs @@ -20,7 +20,7 @@ use arrow::array::ArrayRef; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion::common::ScalarValue; use datafusion::physical_plan::ColumnarValue; -use datafusion_comet_spark_expr::spark_regexp_extract; +use datafusion_comet_spark_expr::{spark_regexp_extract, PatternCache}; use std::hint::black_box; use std::sync::Arc; @@ -47,7 +47,8 @@ fn criterion_benchmark(c: &mut Criterion) { ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"(\d+)-(\d+)".to_string()))), ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), ]; - b.iter(|| black_box(spark_regexp_extract(black_box(&args)).unwrap())) + let cache = PatternCache::new(); + b.iter(|| black_box(spark_regexp_extract(black_box(&args), &cache).unwrap())) }); // Extract the whole match (group 0). @@ -57,7 +58,8 @@ fn criterion_benchmark(c: &mut Criterion) { ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"(\d+)-(\d+)".to_string()))), ColumnarValue::Scalar(ScalarValue::Int32(Some(0))), ]; - b.iter(|| black_box(spark_regexp_extract(black_box(&args)).unwrap())) + let cache = PatternCache::new(); + b.iter(|| black_box(spark_regexp_extract(black_box(&args), &cache).unwrap())) }); // Extract the second capture group. @@ -67,7 +69,8 @@ fn criterion_benchmark(c: &mut Criterion) { ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"(\d+)-(\d+)".to_string()))), ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), ]; - b.iter(|| black_box(spark_regexp_extract(black_box(&args)).unwrap())) + let cache = PatternCache::new(); + b.iter(|| black_box(spark_regexp_extract(black_box(&args), &cache).unwrap())) }); } diff --git a/native/spark-expr/benches/regexp_extract_all.rs b/native/spark-expr/benches/regexp_extract_all.rs index b23b0ac6bca..a2421814c2b 100644 --- a/native/spark-expr/benches/regexp_extract_all.rs +++ b/native/spark-expr/benches/regexp_extract_all.rs @@ -18,7 +18,7 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::common::ScalarValue; use datafusion::physical_plan::ColumnarValue; -use datafusion_comet_spark_expr::spark_regexp_extract_all; +use datafusion_comet_spark_expr::{spark_regexp_extract_all, PatternCache}; use std::hint::black_box; #[path = "common/mod.rs"] @@ -41,7 +41,11 @@ fn criterion_benchmark(c: &mut Criterion) { group.bench_with_input( BenchmarkId::from_parameter(format!("{rows}/{tag}")), &args, - |b, args| b.iter(|| black_box(spark_regexp_extract_all(black_box(args)).unwrap())), + |b, args| { + // One cache per benchmark input mirrors one cache per planned expression. + let cache = PatternCache::new(); + b.iter(|| black_box(spark_regexp_extract_all(black_box(args), &cache).unwrap())) + }, ); } } diff --git a/native/spark-expr/benches/split.rs b/native/spark-expr/benches/split.rs index a28ed394db6..7e39c369ce2 100644 --- a/native/spark-expr/benches/split.rs +++ b/native/spark-expr/benches/split.rs @@ -18,7 +18,7 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::common::ScalarValue; use datafusion::physical_plan::ColumnarValue; -use datafusion_comet_spark_expr::{spark_split, spark_split_sql}; +use datafusion_comet_spark_expr::{spark_split, spark_split_sql, PatternCache}; use std::hint::black_box; #[path = "common/mod.rs"] @@ -40,7 +40,11 @@ fn criterion_benchmark(c: &mut Criterion) { split_group.bench_with_input( BenchmarkId::from_parameter(format!("{rows}/{tag}")), &args, - |b, args| b.iter(|| black_box(spark_split(black_box(args)).unwrap())), + |b, args| { + // One cache per benchmark input mirrors one cache per planned expression. + let cache = PatternCache::new(); + b.iter(|| black_box(spark_split(black_box(args), &cache).unwrap())) + }, ); } } diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index c68b998e617..4b8746ce9bf 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -222,7 +222,11 @@ pub fn create_comet_physical_fun_with_eval_mode( make_comet_scalar_udf!("unbase64", func, without data_type) } "split" => { - let func = Arc::new(crate::string_funcs::spark_split); + // One cache per planned expression: the pattern is a literal, so the regex + // compiles on the first batch and is reused for the rest. + let cache = crate::string_funcs::PatternCache::new(); + let func: ScalarFunctionImplementation = + Arc::new(move |args| crate::string_funcs::spark_split(args, &cache)); make_comet_scalar_udf!("split", func, without data_type) } "split_sql" => { @@ -230,11 +234,15 @@ pub fn create_comet_physical_fun_with_eval_mode( make_comet_scalar_udf!("split_sql", func, without data_type) } "regexp_extract" => { - let func = Arc::new(crate::string_funcs::spark_regexp_extract); + let cache = crate::string_funcs::PatternCache::new(); + let func: ScalarFunctionImplementation = + Arc::new(move |args| crate::string_funcs::spark_regexp_extract(args, &cache)); make_comet_scalar_udf!("regexp_extract", func, without data_type) } "regexp_extract_all" => { - let func = Arc::new(crate::string_funcs::spark_regexp_extract_all); + let cache = crate::string_funcs::PatternCache::new(); + let func: ScalarFunctionImplementation = + Arc::new(move |args| crate::string_funcs::spark_regexp_extract_all(args, &cache)); make_comet_scalar_udf!("regexp_extract_all", func, without data_type) } "get_json_object" => { diff --git a/native/spark-expr/src/string_funcs/mod.rs b/native/spark-expr/src/string_funcs/mod.rs index dc51cfea1b3..996f6bf4b21 100644 --- a/native/spark-expr/src/string_funcs/mod.rs +++ b/native/spark-expr/src/string_funcs/mod.rs @@ -19,6 +19,7 @@ mod base64; mod contains; mod get_json_object; mod levenshtein; +mod pattern_cache; mod regexp_extract; mod regexp_extract_all; mod regexp_extract_common; @@ -29,6 +30,7 @@ pub use base64::spark_base64; pub use contains::SparkContains; pub use get_json_object::spark_get_json_object; pub use levenshtein::spark_levenshtein; +pub use pattern_cache::PatternCache; pub use regexp_extract::spark_regexp_extract; pub use regexp_extract_all::spark_regexp_extract_all; pub use split::{spark_split, spark_split_sql}; diff --git a/native/spark-expr/src/string_funcs/pattern_cache.rs b/native/spark-expr/src/string_funcs/pattern_cache.rs new file mode 100644 index 00000000000..34236640f01 --- /dev/null +++ b/native/spark-expr/src/string_funcs/pattern_cache.rs @@ -0,0 +1,106 @@ +// 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 regex::Regex; +#[cfg(test)] +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Mutex, PoisonError}; + +/// Per-expression cache for a compiled regex. The regexp scalar functions receive their +/// pattern as a per-batch scalar argument even though the serde only plans them with literal +/// patterns, so without this every batch would pay a full regex compile. One slot is enough: +/// a given expression instance sees a single pattern for the lifetime of its plan. +pub struct PatternCache { + cached: Mutex>, + #[cfg(test)] + compile_count: AtomicUsize, +} + +impl PatternCache { + pub fn new() -> Self { + Self { + cached: Mutex::new(None), + #[cfg(test)] + compile_count: AtomicUsize::new(0), + } + } + + /// Return the compiled regex for `pattern`, compiling and caching it only when the + /// pattern differs from the previously cached one. `Regex` clones share the compiled + /// program, so handing out clones is cheap. + pub fn get_or_compile(&self, pattern: &str) -> Result { + // A poisoned lock only means another thread panicked mid-update; the slot is either + // intact or about to be refilled, so recover rather than propagate the panic. + let mut slot = self.cached.lock().unwrap_or_else(PoisonError::into_inner); + if let Some((cached_pattern, regex)) = slot.as_ref() { + if cached_pattern == pattern { + return Ok(regex.clone()); + } + } + #[cfg(test)] + self.compile_count.fetch_add(1, Ordering::Relaxed); + let regex = Regex::new(pattern)?; + *slot = Some((pattern.to_string(), regex.clone())); + Ok(regex) + } + + /// Number of times a regex was actually compiled, for asserting the cache works. + #[cfg(test)] + pub(crate) fn compile_count(&self) -> usize { + self.compile_count.load(Ordering::Relaxed) + } +} + +impl Default for PatternCache { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compiles_once_for_repeated_pattern() { + let cache = PatternCache::new(); + for _ in 0..5 { + let re = cache.get_or_compile(r"(\d+)-(\d+)").unwrap(); + assert!(re.is_match("12-34")); + } + assert_eq!(cache.compile_count(), 1); + } + + #[test] + fn recompiles_when_pattern_changes() { + let cache = PatternCache::new(); + cache.get_or_compile(r"\d+").unwrap(); + cache.get_or_compile(r"[a-z]+").unwrap(); + // Switching back replaces the single slot again. + cache.get_or_compile(r"\d+").unwrap(); + assert_eq!(cache.compile_count(), 3); + } + + #[test] + fn invalid_pattern_errors_and_is_not_cached() { + let cache = PatternCache::new(); + assert!(cache.get_or_compile(r"(unclosed").is_err()); + // A later valid pattern still works. + let re = cache.get_or_compile(r"ok").unwrap(); + assert!(re.is_match("ok")); + } +} diff --git a/native/spark-expr/src/string_funcs/regexp_extract.rs b/native/spark-expr/src/string_funcs/regexp_extract.rs index 1f99b2d5bce..9dd10a182a2 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract.rs @@ -26,6 +26,7 @@ use datafusion::logical_expr::ColumnarValue; use regex::Regex; use std::sync::Arc; +use super::pattern_cache::PatternCache; use super::regexp_extract_common::{parse_args, ParsedArgs}; /// Spark-compatible `regexp_extract(subject, pattern, idx)`. @@ -37,8 +38,11 @@ use super::regexp_extract_common::{parse_args, ParsedArgs}; /// /// Note: this uses the Rust `regex` crate, whose syntax differs from Java's regex engine in /// some ways. The expression is therefore reported as Incompatible. -pub fn spark_regexp_extract(args: &[ColumnarValue]) -> DataFusionResult { - let (regex, group_idx, subject) = match parse_args("regexp_extract", args)? { +pub fn spark_regexp_extract( + args: &[ColumnarValue], + regex_cache: &PatternCache, +) -> DataFusionResult { + let (regex, group_idx, subject) = match parse_args("regexp_extract", args, regex_cache)? { ParsedArgs::Parsed { regex, group_idx, @@ -132,8 +136,12 @@ mod tests { use arrow::array::{LargeStringArray, StringArray}; use datafusion::common::DataFusionError; + fn call_raw(args: &[ColumnarValue]) -> DataFusionResult { + spark_regexp_extract(args, &PatternCache::new()) + } + fn run(args: Vec) -> DataFusionResult>> { - let result = spark_regexp_extract(&args)?; + let result = call_raw(&args)?; match result { ColumnarValue::Array(arr) => { let s = arr @@ -240,7 +248,7 @@ mod tests { #[test] fn group_index_out_of_range_errors() { - let err = spark_regexp_extract(&[array(vec![Some("abc")]), pattern(r"(a)(b)"), idx(3)]) + let err = call_raw(&[array(vec![Some("abc")]), pattern(r"(a)(b)"), idx(3)]) .err() .unwrap(); let msg = err.to_string(); @@ -250,7 +258,7 @@ mod tests { #[test] fn negative_index_errors() { - let err = spark_regexp_extract(&[array(vec![Some("abc")]), pattern(r"(a)"), idx(-1)]) + let err = call_raw(&[array(vec![Some("abc")]), pattern(r"(a)"), idx(-1)]) .err() .unwrap(); let msg = err.to_string(); @@ -260,12 +268,36 @@ mod tests { #[test] fn invalid_regex_errors() { - let err = spark_regexp_extract(&[array(vec![Some("abc")]), pattern(r"(unclosed"), idx(0)]) + let err = call_raw(&[array(vec![Some("abc")]), pattern(r"(unclosed"), idx(0)]) .err() .unwrap(); assert!(err.to_string().contains("`regexp`")); } + /// One expression evaluates many batches; the pattern must compile once and results + /// must stay correct on every batch. + #[test] + fn compiles_regex_once_across_batches() { + let cache = PatternCache::new(); + for batch in 0..4 { + let subject = format!("{batch}00-{batch}99"); + let expected = format!("{batch}00"); + let result = spark_regexp_extract( + &[array(vec![Some(&subject)]), pattern(r"(\d+)-(\d+)"), idx(1)], + &cache, + ) + .unwrap(); + match result { + ColumnarValue::Array(arr) => { + let s = arr.as_any().downcast_ref::().unwrap(); + assert_eq!(s.value(0), expected); + } + other => panic!("unexpected result: {other:?}"), + } + } + assert_eq!(cache.compile_count(), 1); + } + /// `LargeUtf8` subject must still produce a `StringArray` (i32 offsets) so the result type /// matches Spark's `RegExpExtract.dataType` = `StringType`. Regression for the bug where /// `extract_array::` used to build a `LargeStringArray` and trip a type mismatch. @@ -276,7 +308,7 @@ mod tests { None, Some("foo-bar"), ]))); - let result = spark_regexp_extract(&[array, pattern(r"(\d+)-(\d+)"), idx(1)]).unwrap(); + let result = call_raw(&[array, pattern(r"(\d+)-(\d+)"), idx(1)]).unwrap(); match result { ColumnarValue::Array(arr) => { arr.as_any() diff --git a/native/spark-expr/src/string_funcs/regexp_extract_all.rs b/native/spark-expr/src/string_funcs/regexp_extract_all.rs index adf56820b88..b8e37645edc 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract_all.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract_all.rs @@ -28,6 +28,7 @@ use datafusion::logical_expr::ColumnarValue; use regex::Regex; use std::sync::Arc; +use super::pattern_cache::PatternCache; use super::regexp_extract_common::{parse_args, ParsedArgs}; /// Spark-compatible `regexp_extract_all(subject, pattern, idx)`. @@ -40,8 +41,11 @@ use super::regexp_extract_common::{parse_args, ParsedArgs}; /// /// Note: this uses the Rust `regex` crate, whose syntax differs from Java's regex engine in /// some ways. The expression is therefore reported as Incompatible. -pub fn spark_regexp_extract_all(args: &[ColumnarValue]) -> DataFusionResult { - let (regex, group_idx, subject) = match parse_args("regexp_extract_all", args)? { +pub fn spark_regexp_extract_all( + args: &[ColumnarValue], + regex_cache: &PatternCache, +) -> DataFusionResult { + let (regex, group_idx, subject) = match parse_args("regexp_extract_all", args, regex_cache)? { ParsedArgs::Parsed { regex, group_idx, @@ -167,8 +171,12 @@ mod tests { use super::*; use arrow::array::{LargeStringArray, StringArray}; + fn call_raw(args: &[ColumnarValue]) -> DataFusionResult { + spark_regexp_extract_all(args, &PatternCache::new()) + } + fn run(args: Vec) -> DataFusionResult>>> { - let result = spark_regexp_extract_all(&args)?; + let result = call_raw(&args)?; let list = match result { ColumnarValue::Array(arr) => arr, ColumnarValue::Scalar(ScalarValue::List(arr)) => arr as ArrayRef, @@ -317,7 +325,7 @@ mod tests { #[test] fn group_index_out_of_range_errors() { - let err = spark_regexp_extract_all(&[array(vec![Some("abc")]), pattern(r"(a)(b)"), idx(3)]) + let err = call_raw(&[array(vec![Some("abc")]), pattern(r"(a)(b)"), idx(3)]) .err() .unwrap(); let msg = err.to_string(); @@ -327,7 +335,7 @@ mod tests { #[test] fn negative_index_errors() { - let err = spark_regexp_extract_all(&[array(vec![Some("abc")]), pattern(r"(a)"), idx(-1)]) + let err = call_raw(&[array(vec![Some("abc")]), pattern(r"(a)"), idx(-1)]) .err() .unwrap(); let msg = err.to_string(); @@ -337,13 +345,38 @@ mod tests { #[test] fn invalid_regex_errors() { - let err = - spark_regexp_extract_all(&[array(vec![Some("abc")]), pattern(r"(unclosed"), idx(0)]) - .err() - .unwrap(); + let err = call_raw(&[array(vec![Some("abc")]), pattern(r"(unclosed"), idx(0)]) + .err() + .unwrap(); assert!(err.to_string().contains("`regexp`")); } + /// One expression evaluates many batches; the pattern must compile once and results + /// must stay correct on every batch. + #[test] + fn compiles_regex_once_across_batches() { + let cache = PatternCache::new(); + for batch in 0..4 { + let subject = format!("{batch}1-{batch}2, {batch}3-{batch}4"); + let result = spark_regexp_extract_all( + &[array(vec![Some(&subject)]), pattern(r"(\d+)-(\d+)"), idx(1)], + &cache, + ) + .unwrap(); + match result { + ColumnarValue::Array(arr) => { + let list = arr.as_any().downcast_ref::().unwrap(); + let inner = list.value(0); + let strs = inner.as_any().downcast_ref::().unwrap(); + assert_eq!(strs.value(0), format!("{batch}1")); + assert_eq!(strs.value(1), format!("{batch}3")); + } + other => panic!("unexpected result: {other:?}"), + } + } + assert_eq!(cache.compile_count(), 1); + } + /// Regression: `LargeUtf8` subject must still produce a `ListArray` whose inner values /// are a `StringArray` (i32 offsets), matching Spark's `RegExpExtractAll.dataType` = /// `ArrayType(StringType)`. @@ -354,7 +387,7 @@ mod tests { None, Some("4 5"), ]))); - let result = spark_regexp_extract_all(&[array, pattern(r"(\d)"), idx(1)]).unwrap(); + let result = call_raw(&[array, pattern(r"(\d)"), idx(1)]).unwrap(); let list = match result { ColumnarValue::Array(arr) => arr, other => panic!("unexpected result: {other:?}"), diff --git a/native/spark-expr/src/string_funcs/regexp_extract_common.rs b/native/spark-expr/src/string_funcs/regexp_extract_common.rs index 5c02a46b5fa..f45a3de5fb5 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract_common.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract_common.rs @@ -20,6 +20,7 @@ //! module centralizes that parsing and the null short-circuit so each UDF is left with //! its own per-row loop. +use super::pattern_cache::PatternCache; use datafusion::common::{exec_err, DataFusionError, Result as DataFusionResult, ScalarValue}; use datafusion::logical_expr::ColumnarValue; use regex::Regex; @@ -44,6 +45,7 @@ pub(super) enum ParsedArgs<'a> { pub(super) fn parse_args<'a>( fn_name: &'static str, args: &'a [ColumnarValue], + regex_cache: &PatternCache, ) -> DataFusionResult> { if args.len() < 2 || args.len() > 3 { return exec_err!( @@ -82,7 +84,9 @@ pub(super) fn parse_args<'a>( } }; - let regex = Regex::new(pattern).map_err(|e| { + // The pattern is a plan-time literal, so the cache makes this compile a one-time cost + // for the expression instead of a per-batch cost. + let regex = regex_cache.get_or_compile(pattern).map_err(|e| { DataFusionError::Execution(format!( "The value of parameter `regexp` in `{fn_name}` is invalid: '{pattern}' ({e})" )) diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs index 967a898440d..64033440caf 100644 --- a/native/spark-expr/src/string_funcs/split.rs +++ b/native/spark-expr/src/string_funcs/split.rs @@ -29,6 +29,8 @@ use datafusion::logical_expr::ColumnarValue; use regex::Regex; use std::sync::Arc; +use super::pattern_cache::PatternCache; + /// Spark-compatible split function /// Splits a string around matches of a regex pattern with optional limit /// @@ -39,7 +41,10 @@ use std::sync::Arc; /// - limit > 0: At most limit-1 splits, array length <= limit /// - limit = 0: As many splits as possible, trailing empty strings removed /// - limit < 0: As many splits as possible, trailing empty strings kept -pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { +pub fn spark_split( + args: &[ColumnarValue], + regex_cache: &PatternCache, +) -> DataFusionResult { if args.len() < 2 || args.len() > 3 { return exec_err!( "split expects 2 or 3 arguments (string, pattern, [limit]), got {}", @@ -76,7 +81,7 @@ pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { } let pattern_str = pattern.as_ref().unwrap(); - split_array(string_array.as_ref(), pattern_str, limit) + split_array(string_array.as_ref(), pattern_str, limit, regex_cache) } (ColumnarValue::Scalar(ScalarValue::Utf8(string)), ColumnarValue::Scalar(pattern_val)) | ( @@ -97,7 +102,7 @@ pub fn spark_split(args: &[ColumnarValue]) -> DataFusionResult { } }; - let result = split_string(string.as_ref().unwrap(), pattern_str, limit)?; + let result = split_string(string.as_ref().unwrap(), pattern_str, limit, regex_cache)?; let string_array = GenericStringArray::::from(result); let list_array = create_list_array(Arc::new(string_array)); @@ -176,9 +181,11 @@ fn split_array( string_array: &dyn arrow::array::Array, pattern: &str, limit: i32, + regex_cache: &PatternCache, ) -> DataFusionResult { - // Compile regex once for the entire array - let regex = Regex::new(pattern).map_err(|e| { + // The pattern is a plan-time literal, so the cache makes this compile a one-time cost + // for the expression instead of a per-batch cost. + let regex = regex_cache.get_or_compile(pattern).map_err(|e| { DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern, e)) })?; @@ -485,8 +492,13 @@ fn push_split_sql_parts( } } -fn split_string(string: &str, pattern: &str, limit: i32) -> DataFusionResult> { - let regex = Regex::new(pattern).map_err(|e| { +fn split_string( + string: &str, + pattern: &str, + limit: i32, + regex_cache: &PatternCache, +) -> DataFusionResult> { + let regex = regex_cache.get_or_compile(pattern).map_err(|e| { DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern, e)) })?; @@ -592,13 +604,21 @@ mod tests { use super::*; use arrow::array::StringArray; + fn run_split(args: &[ColumnarValue]) -> DataFusionResult { + spark_split(args, &PatternCache::new()) + } + + fn run_split_string(string: &str, pattern: &str, limit: i32) -> DataFusionResult> { + split_string(string, pattern, limit, &PatternCache::new()) + } + #[test] fn test_split_basic() { let string_array = Arc::new(StringArray::from(vec!["a,b,c", "x,y,z"])) as ArrayRef; let pattern = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))); let args = vec![ColumnarValue::Array(string_array), pattern]; - let result = spark_split(&args).unwrap(); + let result = run_split(&args).unwrap(); // Should produce [["a", "b", "c"], ["x", "y", "z"]] assert!(matches!(result, ColumnarValue::Array(_))); } @@ -610,32 +630,32 @@ mod tests { let limit = ColumnarValue::Scalar(ScalarValue::Int32(Some(2))); let args = vec![ColumnarValue::Array(string_array), pattern, limit]; - let result = spark_split(&args).unwrap(); + let result = run_split(&args).unwrap(); // Should produce [["a", "b,c,d"]] assert!(matches!(result, ColumnarValue::Array(_))); } #[test] fn test_split_regex() { - let parts = split_string("foo123bar456baz", r"\d+", -1).unwrap(); + let parts = run_split_string("foo123bar456baz", r"\d+", -1).unwrap(); assert_eq!(parts, vec!["foo", "bar", "baz"]); } #[test] fn test_split_limit_positive() { - let parts = split_string("a,b,c,d,e", ",", 3).unwrap(); + let parts = run_split_string("a,b,c,d,e", ",", 3).unwrap(); assert_eq!(parts, vec!["a", "b", "c,d,e"]); } #[test] fn test_split_limit_zero() { - let parts = split_string("a,b,c,,", ",", 0).unwrap(); + let parts = run_split_string("a,b,c,,", ",", 0).unwrap(); assert_eq!(parts, vec!["a", "b", "c"]); } #[test] fn test_split_limit_negative() { - let parts = split_string("a,b,c,,", ",", -1).unwrap(); + let parts = run_split_string("a,b,c,,", ",", -1).unwrap(); assert_eq!(parts, vec!["a", "b", "c", "", ""]); } @@ -651,7 +671,7 @@ mod tests { let pattern = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))); let args = vec![ColumnarValue::Array(string_array), pattern]; - let result = spark_split(&args).unwrap(); + let result = run_split(&args).unwrap(); match result { ColumnarValue::Array(arr) => { let list_array = arr.as_any().downcast_ref::().unwrap(); @@ -672,7 +692,7 @@ mod tests { #[test] fn test_split_empty_string() { // Test that empty string input produces array with single empty string - let parts = split_string("", ",", -1).unwrap(); + let parts = run_split_string("", ",", -1).unwrap(); assert_eq!(parts, vec![""]); } @@ -754,6 +774,42 @@ mod tests { } } + /// One expression evaluates many batches; the pattern must compile once and results + /// must stay correct on every batch. + #[test] + fn test_split_compiles_regex_once_across_batches() { + let cache = PatternCache::new(); + let pattern = ColumnarValue::Scalar(ScalarValue::Utf8(Some(r"\d+".to_string()))); + + for batch in 0..4 { + let input = format!("foo{batch}bar{batch}baz"); + let string_array = Arc::new(StringArray::from(vec![input.as_str()])) as ArrayRef; + let args = vec![ColumnarValue::Array(string_array), pattern.clone()]; + + let result = spark_split(&args, &cache).unwrap(); + match result { + ColumnarValue::Array(arr) => { + let list_array = arr.as_any().downcast_ref::().unwrap(); + assert_list_value(list_array, 0, &["foo", "bar", "baz"]); + } + _ => panic!("Expected Array result"), + } + } + + assert_eq!(cache.compile_count(), 1); + } + + #[test] + fn test_split_invalid_pattern_errors_at_evaluation() { + let string_array = Arc::new(StringArray::from(vec!["abc"])) as ArrayRef; + let pattern = ColumnarValue::Scalar(ScalarValue::Utf8(Some("(unclosed".to_string()))); + let args = vec![ColumnarValue::Array(string_array), pattern]; + + let err = run_split(&args).err().unwrap(); + let msg = err.to_string(); + assert!(msg.contains("Invalid regex pattern '(unclosed'"), "{msg}"); + } + fn assert_list_value(list_array: &ListArray, row: usize, expected: &[&str]) { let value = list_array.value(row); let strings = value.as_any().downcast_ref::().unwrap(); From f35bc97faf67fa96b2455312dd40a64116e1790c Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 2 Sep 2026 21:26:24 +0700 Subject: [PATCH 2/3] perf: reuse capture locations in regexp_extract_all With the pattern cache handing every invocation a clone of one compiled regex, captures_iter became a bottleneck under concurrent evaluation of the same expression (a sort key evaluated by parallel sort streams): each per-match Captures clones the program's shared group-info Arc, and that refcount turns into a contended cache line. Drive iteration with find_iter, which yields plain spans with identical semantics, and resolve groups through one reused CaptureLocations per batch, matching what regexp_extract already does. This removes the contention and the per-match allocations. --- .../src/string_funcs/regexp_extract_all.rs | 69 +++++++++++++++++-- 1 file changed, 63 insertions(+), 6 deletions(-) diff --git a/native/spark-expr/src/string_funcs/regexp_extract_all.rs b/native/spark-expr/src/string_funcs/regexp_extract_all.rs index b8e37645edc..f6265368278 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract_all.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract_all.rs @@ -102,13 +102,28 @@ fn extract_all_array( let mut null_buffer = BooleanBufferBuilder::new(array.len()); offsets.push(0); + // Reuse one set of capture locations for the whole batch instead of iterating with + // `captures_iter`, which builds a fresh `Captures` per match. Each of those clones the + // compiled regex's shared group-info Arc, and when a sort evaluates the same expression + // from several threads at once that refcount becomes a contended cache line. `find_iter` + // yields plain spans with the exact same iteration semantics, so groups are resolved by + // rerunning the capture engine at each match start into the preallocated buffer. + let mut locations = regex.capture_locations(); for i in 0..array.len() { if array.is_null(i) { offsets.push(values_builder.len() as i32); null_buffer.append(false); } else { - for caps in regex.captures_iter(array.value(i)) { - let s = caps.get(group_idx).map(|m| m.as_str()).unwrap_or(""); + let value = array.value(i); + for m in regex.find_iter(value) { + let matched = regex.captures_read_at(&mut locations, value, m.start()); + debug_assert_eq!( + matched.map(|m| (m.start(), m.end())), + Some((m.start(), m.end())) + ); + let s = locations + .get(group_idx) + .map_or("", |(start, end)| &value[start..end]); values_builder.append_value(s); } offsets.push(values_builder.len() as i32); @@ -128,11 +143,18 @@ fn extract_all_array( } fn extract_one(input: &str, regex: &Regex, group_idx: usize) -> Vec { + let mut locations = regex.capture_locations(); regex - .captures_iter(input) - .map(|caps| { - caps.get(group_idx) - .map(|m| m.as_str().to_string()) + .find_iter(input) + .map(|m| { + let matched = regex.captures_read_at(&mut locations, input, m.start()); + debug_assert_eq!( + matched.map(|m| (m.start(), m.end())), + Some((m.start(), m.end())) + ); + locations + .get(group_idx) + .map(|(start, end)| input[start..end].to_string()) .unwrap_or_default() }) .collect() @@ -312,6 +334,41 @@ mod tests { assert_eq!(result, vec![None]); } + #[test] + fn empty_matches_are_kept_and_iteration_terminates() { + // The regex crate yields empty matches but skips one that sits at the end of the + // previous match, so `a*` on "ba" is ["", "a"] with no trailing empty. + let result = run(vec![array(vec![Some("ba")]), pattern(r"a*"), idx(0)]).unwrap(); + assert_eq!(result, vec![Some(vec![String::new(), "a".to_string()])]); + } + + #[test] + fn empty_matches_advance_over_multibyte_chars() { + let result = run(vec![array(vec![Some("日x本")]), pattern(r"x*"), idx(0)]).unwrap(); + assert_eq!( + result, + vec![Some(vec![String::new(), "x".to_string(), String::new()])] + ); + } + + #[test] + fn anchors_and_word_boundaries_see_full_context() { + let result = run(vec![ + array(vec![Some("cat hat bat")]), + pattern(r"\b\w+\b"), + idx(0), + ]) + .unwrap(); + assert_eq!( + result, + vec![Some(vec![ + "cat".to_string(), + "hat".to_string(), + "bat".to_string() + ])] + ); + } + #[test] fn unmatched_optional_group_returns_empty_string() { let result = run(vec![ From ea8c0e369bad91914703282955cc23e1ae9f93b6 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 4 Sep 2026 12:40:57 +0700 Subject: [PATCH 3/3] perf: search each regexp_extract_all match once Drive the match walk with captures_read_at into one reused CaptureLocations instead of find_iter followed by a second capture search per match. The walk follows the regex crate's iterator rule for empty matches, so results are unchanged, and an equivalence test checks it against captures_iter over empty-match patterns, multibyte haystacks, and out-of-range groups. The benchmark gains the short-row and 8 KB-row cases that exposed the double search. --- .../spark-expr/benches/regexp_extract_all.rs | 42 ++++++ .../src/string_funcs/regexp_extract_all.rs | 134 +++++++++++++----- 2 files changed, 143 insertions(+), 33 deletions(-) diff --git a/native/spark-expr/benches/regexp_extract_all.rs b/native/spark-expr/benches/regexp_extract_all.rs index a2421814c2b..62eb0e38875 100644 --- a/native/spark-expr/benches/regexp_extract_all.rs +++ b/native/spark-expr/benches/regexp_extract_all.rs @@ -29,6 +29,17 @@ const INPUT: &str = "datafusion has datafusion-python, datafusion-comet, datafusion-java as sub projects"; const PATTERN: &str = r"(\w+)-(\w+)"; +/// Short rows with several matches each, so the per-match cost dominates. +const DIGITS_INPUT: &str = "123-456-789-123"; +const DIGITS_PATTERN: &str = r"(\d+)"; +const DIGITS_ROWS: usize = 8_192; + +/// 8 KB rows made of short runs of `a` separated by a single `b`, so a single row carries +/// thousands of matches and the cost of walking a long haystack dominates. +const LONG_ROW_BYTES: usize = 8_192; +const LONG_PATTERN: &str = r"(a+)"; +const LONG_ROWS: usize = 512; + fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("spark_regexp_extract_all"); for rows in ROW_COUNTS { @@ -49,6 +60,37 @@ fn criterion_benchmark(c: &mut Criterion) { ); } } + + let digits_args = vec![ + ColumnarValue::Array(string_array(DIGITS_ROWS, 0.0, |_| DIGITS_INPUT.to_string())), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(DIGITS_PATTERN.to_string()))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ]; + group.bench_with_input( + BenchmarkId::from_parameter(format!("digits/{DIGITS_ROWS}")), + &digits_args, + |b, args| { + // One cache per benchmark input mirrors one cache per planned expression. + let cache = PatternCache::new(); + b.iter(|| black_box(spark_regexp_extract_all(black_box(args), &cache).unwrap())) + }, + ); + + let long_row = "aaab".repeat(LONG_ROW_BYTES / 4); + let long_args = vec![ + ColumnarValue::Array(string_array(LONG_ROWS, 0.0, |_| long_row.clone())), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(LONG_PATTERN.to_string()))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ]; + group.bench_with_input( + BenchmarkId::from_parameter(format!("long_8kb/{LONG_ROWS}")), + &long_args, + |b, args| { + // One cache per benchmark input mirrors one cache per planned expression. + let cache = PatternCache::new(); + b.iter(|| black_box(spark_regexp_extract_all(black_box(args), &cache).unwrap())) + }, + ); group.finish(); } diff --git a/native/spark-expr/src/string_funcs/regexp_extract_all.rs b/native/spark-expr/src/string_funcs/regexp_extract_all.rs index f6265368278..597fc5361ec 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract_all.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract_all.rs @@ -25,7 +25,7 @@ use datafusion::common::{ cast::as_generic_string_array, exec_err, Result as DataFusionResult, ScalarValue, }; use datafusion::logical_expr::ColumnarValue; -use regex::Regex; +use regex::{CaptureLocations, Regex}; use std::sync::Arc; use super::pattern_cache::PatternCache; @@ -102,30 +102,17 @@ fn extract_all_array( let mut null_buffer = BooleanBufferBuilder::new(array.len()); offsets.push(0); - // Reuse one set of capture locations for the whole batch instead of iterating with - // `captures_iter`, which builds a fresh `Captures` per match. Each of those clones the - // compiled regex's shared group-info Arc, and when a sort evaluates the same expression - // from several threads at once that refcount becomes a contended cache line. `find_iter` - // yields plain spans with the exact same iteration semantics, so groups are resolved by - // rerunning the capture engine at each match start into the preallocated buffer. + // One set of capture locations serves the whole batch, so no per-match allocation + // touches the compiled regex's shared group-info Arc from several threads at once. let mut locations = regex.capture_locations(); for i in 0..array.len() { if array.is_null(i) { offsets.push(values_builder.len() as i32); null_buffer.append(false); } else { - let value = array.value(i); - for m in regex.find_iter(value) { - let matched = regex.captures_read_at(&mut locations, value, m.start()); - debug_assert_eq!( - matched.map(|m| (m.start(), m.end())), - Some((m.start(), m.end())) - ); - let s = locations - .get(group_idx) - .map_or("", |(start, end)| &value[start..end]); - values_builder.append_value(s); - } + for_each_group_match(regex, &mut locations, array.value(i), group_idx, |s| { + values_builder.append_value(s) + }); offsets.push(values_builder.len() as i32); null_buffer.append(true); } @@ -144,20 +131,47 @@ fn extract_all_array( fn extract_one(input: &str, regex: &Regex, group_idx: usize) -> Vec { let mut locations = regex.capture_locations(); - regex - .find_iter(input) - .map(|m| { - let matched = regex.captures_read_at(&mut locations, input, m.start()); - debug_assert_eq!( - matched.map(|m| (m.start(), m.end())), - Some((m.start(), m.end())) - ); - locations - .get(group_idx) - .map(|(start, end)| input[start..end].to_string()) - .unwrap_or_default() - }) - .collect() + let mut matches = Vec::new(); + for_each_group_match(regex, &mut locations, input, group_idx, |s| { + matches.push(s.to_string()) + }); + matches +} + +/// Calls `f` with the text of group `group_idx` for every non-overlapping match of `regex` +/// in `haystack`, in the order `captures_iter` yields them, or with the empty string when +/// the group does not participate. Each match costs a single capture search into +/// `locations` and no allocation. +fn for_each_group_match( + regex: &Regex, + locations: &mut CaptureLocations, + haystack: &str, + group_idx: usize, + mut f: impl FnMut(&str), +) { + let mut start = 0; + let mut last_end = None; + while start <= haystack.len() { + let Some(m) = regex.captures_read_at(locations, haystack, start) else { + break; + }; + // The regex crate's iterators drop an empty match that sits at the end of the + // previous match and search again one byte further on. Such a match can only be + // found from that end, so the retry never skips twice in a row. + if m.is_empty() && Some(m.end()) == last_end { + debug_assert_eq!(Some(start), last_end); + start += 1; + continue; + } + start = m.end(); + last_end = Some(m.end()); + let group = locations + .get(group_idx) + .map_or("", |(group_start, group_end)| { + &haystack[group_start..group_end] + }); + f(group); + } } fn null_result(len: Option) -> ColumnarValue { @@ -434,6 +448,60 @@ mod tests { assert_eq!(cache.compile_count(), 1); } + /// The match walk must visit exactly the matches `captures_iter` yields, including the + /// empty-match and multibyte cases where the crate's iterator skips or nudges forward. + #[test] + fn matches_agree_with_captures_iter() { + let patterns = [ + r"a*", + r"(a*)", + r"\b", + r"(\d*)", + r"(?:)", + r"(\d+)", + r"(a+)", + r"zzz", + r"(foo)(bar)?", + r"x*", + r"\b\w+\b", + r"(.)(.)?", + ]; + let haystacks = [ + "", + "a", + "ba", + "aaa", + "123-456-789-123", + "日x本", + "café résumé", + "こんにちは世界", + "a😀b😀", + "foo foo bar", + "cat hat bat", + "a b\tc", + ]; + for pattern in patterns { + let regex = Regex::new(pattern).unwrap(); + let mut locations = regex.capture_locations(); + for haystack in haystacks { + for group_idx in [0usize, 1, 2, 7] { + let expected: Vec = regex + .captures_iter(haystack) + .map(|caps| caps.get(group_idx).map_or("", |m| m.as_str()).to_string()) + .collect(); + let mut actual = Vec::new(); + for_each_group_match(®ex, &mut locations, haystack, group_idx, |s| { + actual.push(s.to_string()) + }); + assert_eq!( + actual, expected, + "pattern {pattern:?} on {haystack:?} group {group_idx}" + ); + } + } + } + } + /// Regression: `LargeUtf8` subject must still produce a `ListArray` whose inner values /// are a `StringArray` (i32 offsets), matching Spark's `RegExpExtractAll.dataType` = /// `ArrayType(StringType)`.