perf(rust): re-use regex capture allocation (#10302)

crates/polars-ops/src/chunked_array/strings/extract.rs

@@ -1,9 +1,11 @@
-use arrow::array::{Array, MutableArray, MutableUtf8Array, StructArray, Utf8Array};
-use polars_arrow::utils::combine_validities_and;
+#[cfg(feature = "extract_groups")]
+use arrow::array::StructArray;
+use arrow::array::{Array, MutableArray, MutableUtf8Array, Utf8Array};
 use polars_core::export::regex::Regex;
 
 use super::*;
 
+#[cfg(feature = "extract_groups")]
 fn extract_groups_array(
     arr: &Utf8Array<i64>,
     reg: &Regex,
@@ -14,51 +16,48 @@ fn extract_groups_array(
         .map(|_| MutableUtf8Array::<i64>::with_capacity(arr.len()))
         .collect::<Vec<_>>();
 
-    arr.into_iter().for_each(|opt_v| {
-        // we combine the null validity later
-        if let Some(value) = opt_v {
-            let caps = reg.captures(value);
-            match caps {
-                None => builders.iter_mut().for_each(|arr| arr.push_null()),
-                Some(caps) => {
-                    caps.iter()
-                        .skip(1) // skip 0th group
-                        .zip(builders.iter_mut())
-                        .for_each(|(m, builder)| builder.push(m.map(|m| m.as_str())))
+    let mut locs = reg.capture_locations();
+    for opt_v in arr {
+        if let Some(s) = opt_v {
+            if reg.captures_read(&mut locs, s).is_some() {
+                for (i, builder) in builders.iter_mut().enumerate() {
+                    builder.push(locs.get(i + 1).map(|(start, stop)| &s[start..stop]));
                 }
+                continue;
             }
         }
-    });
+
+        // Push nulls if either the string is null or there was no match. We
+        // distinguish later between the two by copying arr's validity mask.
+        builders.iter_mut().for_each(|arr| arr.push_null());
+    }
 
     let values = builders
         .into_iter()
-        .map(|group_array| {
-            let group_array: Utf8Array<i64> = group_array.into();
-            let final_validity = combine_validities_and(group_array.validity(), arr.validity());
-            group_array.with_validity(final_validity).to_boxed()
+        .map(|a| {
+            let immutable_a: Utf8Array<i64> = a.into();
+            immutable_a.to_boxed()
         })
         .collect();
-
-    Ok(StructArray::new(data_type.clone(), values, None).boxed())
+    Ok(StructArray::new(data_type.clone(), values, arr.validity().cloned()).boxed())
 }
 
+#[cfg(feature = "extract_groups")]
 pub(super) fn extract_groups(
     ca: &Utf8Chunked,
     pat: &str,
     dtype: &DataType,
 ) -> PolarsResult<Series> {
     let reg = Regex::new(pat)?;
     let n_fields = reg.captures_len();
-
     if n_fields == 1 {
         return StructChunked::new(ca.name(), &[Series::new_null(ca.name(), ca.len())])
             .map(|ca| ca.into_series());
     }
 
     let data_type = dtype.to_arrow();
-    // impl error if it isn't a struct
     let DataType::Struct(fields) = dtype else {
-        unreachable!()
+        unreachable!() // Implementation error if it isn't a struct.
     };
     let names = fields
         .iter()
@@ -72,3 +71,42 @@ pub(super) fn extract_groups(
 
     Series::try_from((ca.name(), chunks))
 }
+
+fn extract_group_array(
+    arr: &Utf8Array<i64>,
+    reg: &Regex,
+    group_index: usize,
+) -> PolarsResult<Utf8Array<i64>> {
+    let mut builder = MutableUtf8Array::<i64>::with_capacity(arr.len());
+
+    let mut locs = reg.capture_locations();
+    for opt_v in arr {
+        if let Some(s) = opt_v {
+            if reg.captures_read(&mut locs, s).is_some() {
+                builder.push(locs.get(group_index).map(|(start, stop)| &s[start..stop]));
+                continue;
+            }
+        }
+
+        // Push null if either the string is null or there was no match.
+        builder.push_null();
+    }
+
+    Ok(builder.into())
+}
+
+pub(super) fn extract_group(
+    ca: &Utf8Chunked,
+    pat: &str,
+    group_index: usize,
+) -> PolarsResult<Utf8Chunked> {
+    let reg = Regex::new(pat)?;
+
+    let chunks = ca
+        .downcast_iter()
+        .map(|array| Ok(extract_group_array(array, &reg, group_index)?.to_boxed()))
+        .collect::<PolarsResult<Vec<Box<dyn Array>>>>()?;
+
+    // SAFETY: all chunks have type Utf8Array
+    unsafe { Ok(ChunkedArray::from_chunks(ca.name(), chunks)) }
+}

crates/polars-ops/src/chunked_array/strings/mod.rs

@@ -1,6 +1,6 @@
 #[cfg(feature = "strings")]
 mod case;
-#[cfg(feature = "extract_groups")]
+#[cfg(feature = "strings")]
 mod extract;
 #[cfg(feature = "extract_jsonpath")]
 mod json_path;

crates/polars-ops/src/chunked_array/strings/namespace.rs

@@ -1,5 +1,3 @@
-use std::borrow::Cow;
-
 #[cfg(feature = "string_encoding")]
 use base64::engine::general_purpose;
 #[cfg(feature = "string_encoding")]
@@ -15,11 +13,6 @@ use super::*;
 #[cfg(feature = "binary_encoding")]
 use crate::chunked_array::binary::BinaryNameSpaceImpl;
 
-fn f_regex_extract<'a>(reg: &Regex, input: &'a str, group_index: usize) -> Option<Cow<'a, str>> {
-    reg.captures(input)
-        .and_then(|cap| cap.get(group_index).map(|m| Cow::Borrowed(m.as_str())))
-}
-
 pub trait Utf8NameSpaceImpl: AsUtf8 {
     #[cfg(not(feature = "binary_encoding"))]
     fn hex_decode(&self) -> PolarsResult<Utf8Chunked> {
@@ -298,8 +291,7 @@ pub trait Utf8NameSpaceImpl: AsUtf8 {
     /// Extract the nth capture group from pattern
     fn extract(&self, pat: &str, group_index: usize) -> PolarsResult<Utf8Chunked> {
         let ca = self.as_utf8();
-        let reg = Regex::new(pat)?;
-        Ok(ca.apply_on_opt(|e| e.and_then(|input| f_regex_extract(&reg, input, group_index))))
+        super::extract::extract_group(ca, pat, group_index)
     }
 
     /// Extract each successive non-overlapping regex match in an individual string as an array