From ed8d133ddebca3fd2cfbf1a8f8787cc573cb8e68 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 11:44:24 -0600 Subject: [PATCH] perf: optimize spark_regexp_extract_all in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + .../spark-expr/benches/regexp_extract_all.rs | 74 ++++++++++++ .../src/string_funcs/regexp_extract_all.rs | 111 +++++++++++++++--- 3 files changed, 170 insertions(+), 19 deletions(-) create mode 100644 native/spark-expr/benches/regexp_extract_all.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..0022dd82b9 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -139,3 +139,7 @@ harness = false [[bench]] name = "to_json" harness = false + +[[bench]] +name = "regexp_extract_all" +harness = false diff --git a/native/spark-expr/benches/regexp_extract_all.rs b/native/spark-expr/benches/regexp_extract_all.rs new file mode 100644 index 0000000000..3e8f8874ef --- /dev/null +++ b/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> = (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 { + 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); 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..e7887a4a23 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract_all.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract_all.rs @@ -93,23 +93,21 @@ fn extract_all_array( 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 = 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( )) } +/// 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(array: &GenericStringArray) -> 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 { - 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(&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 = 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) -> ColumnarValue {