improve performance of regexp_count

Author: Dima

datafusion/functions/src/regex/regexpcount.rs

@@ -30,7 +30,6 @@ use datafusion_expr::{
 };
 use itertools::izip;
 use regex::Regex;
-use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::sync::{Arc, OnceLock};
 
@@ -310,12 +309,13 @@ where
                 Some(regex) => regex,
             };
 
-            let pattern = compile_regex(regex, flags_scalar)?;
+            let pattern = get_pattern(regex, flags_scalar)?;
+            let re = compile_regex(pattern)?;
 
             Ok(Arc::new(Int64Array::from_iter_values(
                 values
                     .iter()
-                    .map(|value| count_matches(value, &pattern, start_scalar))
+                    .map(|value| count_matches(value, &re, start_scalar))
                     .collect::<Result<Vec<i64>, ArrowError>>()?,
             )))
         }
@@ -356,15 +356,15 @@ where
                 Some(regex) => regex,
             };
 
-            let pattern = compile_regex(regex, flags_scalar)?;
-
+            let pattern = get_pattern(regex, flags_scalar)?;
+            let re = compile_regex(pattern)?;
             let start_array = start_array.unwrap();
 
             Ok(Arc::new(Int64Array::from_iter_values(
                 values
                     .iter()
                     .zip(start_array.iter())
-                    .map(|(value, start)| count_matches(value, &pattern, start))
+                    .map(|(value, start)| count_matches(value, &re, start))
                     .collect::<Result<Vec<i64>, ArrowError>>()?,
             )))
         }
@@ -549,34 +549,37 @@ where
     }
 }
 
-fn compile_and_cache_regex(
-    regex: &str,
-    flags: Option<&str>,
-    regex_cache: &mut HashMap<String, Regex>,
-) -> Result<Regex, ArrowError> {
-    match regex_cache.entry(regex.to_string()) {
-        Entry::Vacant(entry) => {
-            let compiled = compile_regex(regex, flags)?;
-            entry.insert(compiled.clone());
-            Ok(compiled)
-        }
-        Entry::Occupied(entry) => Ok(entry.get().to_owned()),
+fn compile_and_cache_regex<'a>(
+    regex: &'a str,
+    flags: Option<&'a str>,
+    regex_cache: &'a mut HashMap<String, Regex>,
+) -> Result<&'a Regex, ArrowError> {
+    let pattern = get_pattern(regex, flags)?;
+
+    if regex_cache.contains_key(&pattern) {
+        return Ok(regex_cache.get(&pattern).unwrap());
     }
+
+    let re = compile_regex(pattern.clone())?;
+    regex_cache.insert(pattern.clone(), re);
+    Ok(regex_cache.get(&pattern).unwrap())
 }
 
-fn compile_regex(regex: &str, flags: Option<&str>) -> Result<Regex, ArrowError> {
-    let pattern = match flags {
-        None | Some("") => regex.to_string(),
+fn get_pattern(regex: &str, flags: Option<&str>) -> Result<String, ArrowError> {
+    match flags {
+        None | Some("") => Ok(regex.to_string()),
         Some(flags) => {
             if flags.contains("g") {
                 return Err(ArrowError::ComputeError(
                     "regexp_count() does not support global flag".to_string(),
                 ));
             }
-            format!("(?{}){}", flags, regex)
+            Ok(format!("(?{}){}", flags, regex))
         }
-    };
+    }
+}
 
+fn compile_regex(pattern: String) -> Result<Regex, ArrowError> {
     Regex::new(&pattern).map_err(|_| {
         ArrowError::ComputeError(format!(
             "Regular expression did not compile: {}",