perf: optimize spark_regexp_extract_all in datafusion-comet-spark-expr

native/spark-expr/Cargo.toml

@@ -139,3 +139,7 @@ harness = false
 [[bench]]
 name = "to_json"
 harness = false
+
+[[bench]]
+name = "regexp_extract_all"
+harness = false

native/spark-expr/benches/regexp_extract_all.rs

@@ -0,0 +1,74 @@
+// 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 arrow::array::StringArray;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::ScalarValue;
+use datafusion::logical_expr::ColumnarValue;
+use datafusion_comet_spark_expr::spark_regexp_extract_all;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const NUM_ROWS: usize = 4096;
+
+fn subjects() -> ColumnarValue {
+    let values: Vec<Option<String>> = (0..NUM_ROWS)
+        .map(|i| match i % 8 {
+            0 => None,
+            1 => Some("no digits at all in this row".to_string()),
+            2 => Some(format!("{}-{}", i, i + 1)),
+            _ => Some(format!(
+                "{}-{}, {}-{} and {}-{}",
+                i,
+                i + 1,
+                i * 3,
+                i * 3 + 7,
+                i * 11,
+                i * 11 + 13
+            )),
+        })
+        .collect();
+    ColumnarValue::Array(Arc::new(StringArray::from(values)))
+}
+
+fn args(subjects: &ColumnarValue, pattern: &str, idx: i32) -> Vec<ColumnarValue> {
+    vec![
+        subjects.clone(),
+        ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern.to_string()))),
+        ColumnarValue::Scalar(ScalarValue::Int32(Some(idx))),
+    ]
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let subjects = subjects();
+    let mut group = c.benchmark_group("regexp_extract_all");
+
+    let capture_group = args(&subjects, r"(\d+)-(\d+)", 1);
+    group.bench_function("capture_group", |b| {
+        b.iter(|| black_box(spark_regexp_extract_all(black_box(&capture_group)).unwrap()))
+    });
+
+    let whole_match = args(&subjects, r"(\d+)-(\d+)", 0);
+    group.bench_function("whole_match", |b| {
+        b.iter(|| black_box(spark_regexp_extract_all(black_box(&whole_match)).unwrap()))
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

native/spark-expr/src/string_funcs/regexp_extract_all.rs

@@ -93,23 +93,21 @@ fn extract_all_array<O: OffsetSizeTrait>(
     regex: &Regex,
     group_idx: usize,
 ) -> ArrayRef {
-    let mut values_builder = StringBuilder::new();
+    let mut values_builder = StringBuilder::with_capacity(array.len(), values_capacity(array));
     let mut offsets: Vec<i32> = Vec::with_capacity(array.len() + 1);
     let mut null_buffer = BooleanBufferBuilder::new(array.len());
     offsets.push(0);
 
-    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);
+    let mut matcher = GroupMatcher::new(regex, group_idx);
+    for value in array.iter() {
+        match value {
+            Some(s) => {
+                matcher.for_each_match(s, |s| values_builder.append_value(s));
+                null_buffer.append(true);
             }
-            offsets.push(values_builder.len() as i32);
-            null_buffer.append(true);
+            None => null_buffer.append(false),
         }
+        offsets.push(values_builder.len() as i32);
     }
 
     let values = Arc::new(values_builder.finish()) as ArrayRef;
@@ -123,15 +121,90 @@ fn extract_all_array<O: OffsetSizeTrait>(
     ))
 }
 
+/// Byte capacity to pre-allocate for the extracted strings of `array`.
+///
+/// The extracted text is a subset of the subject, so the subject's own byte length is an upper
+/// bound; it is capped so that a large batch does not reserve far more than it is likely to use.
+fn values_capacity<O: OffsetSizeTrait>(array: &GenericStringArray<O>) -> usize {
+    const MAX_CAPACITY: usize = 1 << 20;
+    let offsets = array.value_offsets();
+    let subject_bytes = offsets[array.len()].as_usize() - offsets[0].as_usize();
+    subject_bytes.min(MAX_CAPACITY)
+}
+
 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 matches = Vec::new();
+    GroupMatcher::new(regex, group_idx).for_each_match(input, |s| matches.push(s.to_string()));
+    matches
+}
+
+/// Walks the non-overlapping matches of a regex, yielding the text of one capture group.
+///
+/// This is equivalent to iterating `Regex::captures_iter` and reading group `group_idx` from
+/// each `Captures`, but `captures_iter` allocates a fresh capture-slot buffer for every match.
+/// Here a single `CaptureLocations` is reused across all matches and all rows, and when only
+/// the whole match is needed (`idx = 0`) the capture groups are not resolved at all.
+struct GroupMatcher<'a> {
+    regex: &'a Regex,
+    group_idx: usize,
+    locations: regex::CaptureLocations,
+}
+
+impl<'a> GroupMatcher<'a> {
+    fn new(regex: &'a Regex, group_idx: usize) -> Self {
+        Self {
+            regex,
+            group_idx,
+            locations: regex.capture_locations(),
+        }
+    }
+
+    /// Calls `emit` once per match with the matched text of the capture group, using the empty
+    /// string for a group that did not participate in the match.
+    fn for_each_match<F: FnMut(&str)>(&mut self, input: &str, mut emit: F) {
+        if self.group_idx == 0 {
+            for m in self.regex.find_iter(input) {
+                emit(m.as_str());
+            }
+            return;
+        }
+
+        let mut search_start = 0;
+        let mut last_match_end: Option<usize> = None;
+        while search_start <= input.len() {
+            let Some(m) = self
+                .regex
+                .captures_read_at(&mut self.locations, input, search_start)
+            else {
+                return;
+            };
+            // The same empty match is never reported twice: an empty match landing on the end of
+            // the previous match is dropped.
+            if m.is_empty() && Some(m.end()) == last_match_end {
+                search_start = next_char_start(input, m.end());
+                continue;
+            }
+            emit(
+                self.locations
+                    .get(self.group_idx)
+                    .map_or("", |(start, end)| &input[start..end]),
+            );
+            last_match_end = Some(m.end());
+            // Resuming at the end of an empty match would just find it again, so step over the
+            // character it sits on.
+            search_start = if m.is_empty() {
+                next_char_start(input, m.end())
+            } else {
+                m.end()
+            };
+        }
+    }
+}
+
+/// Byte index of the character following the one starting at `index`. An `index` at the end of
+/// `input` yields an index past the end, which terminates the search.
+fn next_char_start(input: &str, index: usize) -> usize {
+    index + input[index..].chars().next().map_or(1, char::len_utf8)
 }
 
 fn null_result(len: Option<usize>) -> ColumnarValue {