perf: compile user regex patterns once per planned expression

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()))
     });
 }
 

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();
 }
 

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()))
+                },
             );
         }
     }

native/spark-expr/src/comet_scalar_funcs.rs

@@ -223,19 +223,27 @@ 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" => {
             let func = Arc::new(crate::string_funcs::spark_split_sql);
             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" => {

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};

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"));
+    }
+}

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()