You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/12 22:37:13 UTC
(datafusion-comet) 01/01: perf: compile user regex patterns once per planned expression (#5612)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5612-0eec0601b6637db1d85fac9b171f31a0df26c5e2
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit e16461e0c6b79fd5da3599e17c19a293673b3136
Author: dustin <dw...@users.noreply.github.com>
AuthorDate: Sat Sep 12 20:34:59 2026 +0000
perf: compile user regex patterns once per planned expression (#5612)
* 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.
* 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.
* 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.
---
native/spark-expr/benches/regexp_extract.rs | 11 +-
native/spark-expr/benches/regexp_extract_all.rs | 50 ++++-
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 +
.../spark-expr/src/string_funcs/pattern_cache.rs | 106 +++++++++++
.../spark-expr/src/string_funcs/regexp_extract.rs | 46 ++++-
.../src/string_funcs/regexp_extract_all.rs | 204 ++++++++++++++++++---
.../src/string_funcs/regexp_extract_common.rs | 6 +-
native/spark-expr/src/string_funcs/split.rs | 86 +++++++--
10 files changed, 476 insertions(+), 57 deletions(-)
diff --git a/native/spark-expr/benches/regexp_extract.rs b/native/spark-expr/benches/regexp_extract.rs
index 31cf9f32a7..9f514116c9 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 b23b0ac6bc..62eb0e3887 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 a28ed394db..7e39c369ce 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 35aa5c6ab3..7f8e27930b 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -234,7 +234,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" => {
@@ -242,11 +246,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 94afc12d05..d406c620a1 100644
--- a/native/spark-expr/src/string_funcs/mod.rs
+++ b/native/spark-expr/src/string_funcs/mod.rs
@@ -20,6 +20,7 @@ mod concat_ws;
mod contains;
mod get_json_object;
mod levenshtein;
+mod pattern_cache;
mod regexp_extract;
mod regexp_extract_all;
mod regexp_extract_common;
@@ -31,6 +32,7 @@ pub use concat_ws::CometConcatWs;
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 0000000000..34236640f0
--- /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<Option<(String, Regex)>>,
+ #[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<Regex, regex::Error> {
+ // 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 1f99b2d5bc..9dd10a182a 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<ColumnarValue> {
- let (regex, group_idx, subject) = match parse_args("regexp_extract", args)? {
+pub fn spark_regexp_extract(
+ args: &[ColumnarValue],
+ regex_cache: &PatternCache,
+) -> DataFusionResult<ColumnarValue> {
+ 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<ColumnarValue> {
+ spark_regexp_extract(args, &PatternCache::new())
+ }
+
fn run(args: Vec<ColumnarValue>) -> DataFusionResult<Vec<Option<String>>> {
- 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::<StringArray>().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::<i64>` 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 adf56820b8..597fc5361e 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<ColumnarValue> {
- let (regex, group_idx, subject) = match parse_args("regexp_extract_all", args)? {
+pub fn spark_regexp_extract_all(
+ args: &[ColumnarValue],
+ regex_cache: &PatternCache,
+) -> DataFusionResult<ColumnarValue> {
+ 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<O: OffsetSizeTrait>(
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<O: OffsetSizeTrait>(
}
fn extract_one(input: &str, regex: &Regex, group_idx: usize) -> Vec<String> {
- 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<usize>) -> ColumnarValue {
@@ -167,8 +207,12 @@ mod tests {
use super::*;
use arrow::array::{LargeStringArray, StringArray};
+ fn call_raw(args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
+ spark_regexp_extract_all(args, &PatternCache::new())
+ }
+
fn run(args: Vec<ColumnarValue>) -> DataFusionResult<Vec<Option<Vec<String>>>> {
- 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::<ListArray>().unwrap();
+ let inner = list.value(0);
+ let strs = inner.as_any().downcast_ref::<StringArray>().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<String> = 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 5c02a46b5f..f45a3de5fb 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<ParsedArgs<'a>> {
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 967a898440..64033440ca 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<ColumnarValue> {
+pub fn spark_split(
+ args: &[ColumnarValue],
+ regex_cache: &PatternCache,
+) -> DataFusionResult<ColumnarValue> {
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<ColumnarValue> {
}
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<ColumnarValue> {
}
};
- 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::<i32>::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<ColumnarValue> {
- // 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<O: OffsetSizeTrait>(
}
}
-fn split_string(string: &str, pattern: &str, limit: i32) -> DataFusionResult<Vec<String>> {
- let regex = Regex::new(pattern).map_err(|e| {
+fn split_string(
+ string: &str,
+ pattern: &str,
+ limit: i32,
+ regex_cache: &PatternCache,
+) -> DataFusionResult<Vec<String>> {
+ 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<ColumnarValue> {
+ spark_split(args, &PatternCache::new())
+ }
+
+ fn run_split_string(string: &str, pattern: &str, limit: i32) -> DataFusionResult<Vec<String>> {
+ 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::<ListArray>().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::<ListArray>().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::<StringArray>().unwrap();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org