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..62eb0e38875 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"] @@ -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 { @@ -41,10 +52,45 @@ 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())) + }, ); } } + + 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/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 5e7562c4180..62779326ecd 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -226,7 +226,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" => { @@ -234,11 +238,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..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,9 +25,10 @@ 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; 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, @@ -98,15 +102,17 @@ fn extract_all_array( let mut null_buffer = BooleanBufferBuilder::new(array.len()); offsets.push(0); + // 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 { - for caps in regex.captures_iter(array.value(i)) { - let s = caps.get(group_idx).map(|m| m.as_str()).unwrap_or(""); - 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); } @@ -124,14 +130,48 @@ fn extract_all_array( } fn extract_one(input: &str, regex: &Regex, group_idx: usize) -> Vec { - regex - .captures_iter(input) - .map(|caps| { - caps.get(group_idx) - .map(|m| m.as_str().to_string()) - .unwrap_or_default() - }) - .collect() + let mut locations = regex.capture_locations(); + 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 { @@ -167,8 +207,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, @@ -304,6 +348,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![ @@ -317,7 +396,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 +406,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 +416,92 @@ 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); + } + + /// 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)`. @@ -354,7 +512,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();