From 50fee655a1f513b1d695c87060d63f08ce31fc84 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 11 Jul 2026 07:50:14 -0600 Subject: [PATCH] perf: optimize spark_regexp_extract in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/regexp_extract.rs | 75 +++++++++++++++++++ .../src/string_funcs/regexp_extract.rs | 16 +++- 3 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 native/spark-expr/benches/regexp_extract.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 4addc026d8..ad267f86a8 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -131,3 +131,7 @@ harness = false [[bench]] name = "array_size" harness = false + +[[bench]] +name = "regexp_extract" +harness = false diff --git a/native/spark-expr/benches/regexp_extract.rs b/native/spark-expr/benches/regexp_extract.rs new file mode 100644 index 0000000000..31cf9f32a7 --- /dev/null +++ b/native/spark-expr/benches/regexp_extract.rs @@ -0,0 +1,75 @@ +// 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::builder::StringBuilder; +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 std::hint::black_box; +use std::sync::Arc; + +fn create_string_array(size: usize) -> ArrayRef { + let mut builder = StringBuilder::new(); + for i in 0..size { + if i % 10 == 0 { + builder.append_null(); + } else { + builder.append_value(format!("{:03}-{:03}", i % 1000, (i * 7) % 1000)); + } + } + Arc::new(builder.finish()) +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + let string_array = create_string_array(size); + + // Extract the first capture group of a two-group pattern (the common case). + c.bench_function("spark_regexp_extract: group 1", |b| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&string_array)), + 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())) + }); + + // Extract the whole match (group 0). + c.bench_function("spark_regexp_extract: group 0", |b| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&string_array)), + 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())) + }); + + // Extract the second capture group. + c.bench_function("spark_regexp_extract: group 2", |b| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&string_array)), + 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())) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/string_funcs/regexp_extract.rs b/native/spark-expr/src/string_funcs/regexp_extract.rs index 156241d6a1..1f99b2d5bc 100644 --- a/native/spark-expr/src/string_funcs/regexp_extract.rs +++ b/native/spark-expr/src/string_funcs/regexp_extract.rs @@ -87,14 +87,22 @@ fn extract_array( group_idx: usize, ) -> ArrayRef { let mut builder = StringBuilder::with_capacity(array.len(), array.value_data().len()); + // Reuse a single set of capture locations across every row. `Regex::captures` + // allocates a fresh location buffer on each call; `captures_read` fills this + // preallocated buffer in place, so the per-row heap allocation is paid only once. + let mut locations = regex.capture_locations(); for i in 0..array.len() { if array.is_null(i) { builder.append_null(); } else { - let extracted = match regex.captures(array.value(i)) { - Some(caps) => caps.get(group_idx).map(|m| m.as_str()).unwrap_or(""), - None => "", - }; + let value = array.value(i); + // `captures_read` returns `Some` iff the pattern matched; `locations.get` + // is `None` when the requested group did not participate in the match, + // matching the previous `caps.get(group_idx)` behavior (empty string). + let extracted = regex + .captures_read(&mut locations, value) + .and_then(|_| locations.get(group_idx)) + .map_or("", |(start, end)| &value[start..end]); builder.append_value(extracted); } }