Skip to content

Commit 17ec083

Browse files
andygrovecomphead
andauthored
perf: optimize regexp_match for literal pattern usage (20% faster) (#23547)
## Which issue does this PR close? N/A ## Rationale for this change Optimize existing expressions ## What changes are included in this PR? Pass a literal regexp pattern/flags to arrow's kernel as scalar Datums so the regex is compiled once per batch, instead of expanding the literal to a full array and forcing a per-row HashMap<String, Regex> cache lookup (plus a per-row format! allocation when flags are present). ## Are these changes tested? Existing tests + new unit tests. Benchmark (criterion): - regexp_match_1000 literal pattern utf8view: 21.394% faster (base 249069ns -> cand 195784ns) - regexp_match_1000 pattern array: 2.996% faster (base 190020ns -> cand 184327ns) - regexp_match_1000 literal pattern: 23.964% faster (base 261795ns -> cand 199059ns) - regexp_match_1000 literal pattern and flags: 37.227% faster (base 280047ns -> cand 175794ns) Full criterion output: ```text regexp_match_1000 literal pattern time: [198.39 µs 198.66 µs 199.03 µs] change: [−24.244% −23.964% −23.600%] (p = 0.00 < 0.05) Performance has improved. Found 9 outliers among 100 measurements (9.00%) 1 (1.00%) low mild 4 (4.00%) high mild 4 (4.00%) high severe regexp_match_1000 literal pattern and flags time: [175.59 µs 175.68 µs 175.81 µs] change: [−37.342% −37.227% −37.124%] (p = 0.00 < 0.05) Performance has improved. Found 2 outliers among 100 measurements (2.00%) 1 (1.00%) high mild 1 (1.00%) high severe regexp_match_1000 literal pattern utf8view time: [195.75 µs 195.85 µs 195.98 µs] change: [−21.511% −21.394% −21.298%] (p = 0.00 < 0.05) Performance has improved. Found 4 outliers among 100 measurements (4.00%) 3 (3.00%) high mild 1 (1.00%) high severe regexp_match_1000 pattern array time: [184.11 µs 184.22 µs 184.37 µs] change: [−3.1007% −2.9962% −2.8581%] (p = 0.00 < 0.05) Performance has improved. Found 5 outliers among 100 measurements (5.00%) 1 (1.00%) high mild 4 (4.00%) high severe ``` ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: Oleks V <comphead@users.noreply.github.com>
1 parent 1e58928 commit 17ec083

3 files changed

Lines changed: 273 additions & 1 deletion

File tree

datafusion/functions/Cargo.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,11 @@ harness = false
163163
name = "to_hex"
164164
required-features = ["string_expressions"]
165165

166+
[[bench]]
167+
harness = false
168+
name = "regexp_match"
169+
required-features = ["regex_expressions"]
170+
166171
[[bench]]
167172
harness = false
168173
name = "regx"
Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Benchmarks `regexp_match` through `invoke_with_args`, which is how a query
19+
//! plan calls it. The pattern (and flags) are literals, as in
20+
//! `regexp_match(col, '[a-z]+')`.
21+
22+
use std::hint::black_box;
23+
use std::sync::Arc;
24+
25+
use arrow::array::{ArrayRef, StringArray};
26+
use arrow::compute::cast;
27+
use arrow::datatypes::{DataType, Field};
28+
use criterion::{Criterion, criterion_group, criterion_main};
29+
use datafusion_common::ScalarValue;
30+
use datafusion_common::config::ConfigOptions;
31+
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
32+
use datafusion_functions::regex::regexpmatch::RegexpMatchFunc;
33+
use rand::Rng;
34+
use rand::distr::Alphanumeric;
35+
use rand::rngs::ThreadRng;
36+
37+
const SIZE: usize = 1000;
38+
const PATTERN: &str = ".*([A-Z]{1}).*";
39+
40+
fn data(rng: &mut ThreadRng) -> StringArray {
41+
(0..SIZE)
42+
.map(|_| {
43+
rng.sample_iter(&Alphanumeric)
44+
.take(7)
45+
.map(char::from)
46+
.collect::<String>()
47+
})
48+
.collect::<Vec<_>>()
49+
.into()
50+
}
51+
52+
fn run(c: &mut Criterion, name: &str, values: &ArrayRef, args: &[ColumnarValue]) {
53+
let func = RegexpMatchFunc::new();
54+
let arg_fields: Vec<_> = args
55+
.iter()
56+
.enumerate()
57+
.map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into())
58+
.collect();
59+
let return_field = Arc::new(Field::new_list(
60+
"f",
61+
Field::new_list_field(values.data_type().clone(), true),
62+
true,
63+
));
64+
let config_options = Arc::new(ConfigOptions::default());
65+
66+
c.bench_function(name, |b| {
67+
b.iter(|| {
68+
black_box(
69+
func.invoke_with_args(ScalarFunctionArgs {
70+
args: args.to_vec(),
71+
arg_fields: arg_fields.clone(),
72+
number_rows: SIZE,
73+
return_field: Arc::clone(&return_field),
74+
config_options: Arc::clone(&config_options),
75+
})
76+
.expect("regexp_match should work on valid values"),
77+
)
78+
})
79+
});
80+
}
81+
82+
fn criterion_benchmark(c: &mut Criterion) {
83+
let mut rng = rand::rng();
84+
let utf8 = Arc::new(data(&mut rng)) as ArrayRef;
85+
let utf8view = cast(&utf8, &DataType::Utf8View).unwrap();
86+
87+
run(
88+
c,
89+
"regexp_match_1000 literal pattern",
90+
&utf8,
91+
&[
92+
ColumnarValue::Array(Arc::clone(&utf8)),
93+
ColumnarValue::Scalar(ScalarValue::Utf8(Some(PATTERN.to_string()))),
94+
],
95+
);
96+
97+
run(
98+
c,
99+
"regexp_match_1000 literal pattern and flags",
100+
&utf8,
101+
&[
102+
ColumnarValue::Array(Arc::clone(&utf8)),
103+
ColumnarValue::Scalar(ScalarValue::Utf8(Some(PATTERN.to_string()))),
104+
ColumnarValue::Scalar(ScalarValue::Utf8(Some("i".to_string()))),
105+
],
106+
);
107+
108+
run(
109+
c,
110+
"regexp_match_1000 literal pattern utf8view",
111+
&utf8view,
112+
&[
113+
ColumnarValue::Array(Arc::clone(&utf8view)),
114+
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(PATTERN.to_string()))),
115+
],
116+
);
117+
118+
// Covers the path where the pattern varies per row and so cannot be
119+
// compiled once for the whole array.
120+
let patterns = Arc::new(StringArray::from(
121+
(0..SIZE)
122+
.map(|i| if i % 2 == 0 { PATTERN } else { "^(A).*" })
123+
.collect::<Vec<_>>(),
124+
)) as ArrayRef;
125+
run(
126+
c,
127+
"regexp_match_1000 pattern array",
128+
&utf8,
129+
&[
130+
ColumnarValue::Array(Arc::clone(&utf8)),
131+
ColumnarValue::Array(patterns),
132+
],
133+
);
134+
}
135+
136+
criterion_group!(benches, criterion_benchmark);
137+
criterion_main!(benches);

datafusion/functions/src/regex/regexpmatch.rs

Lines changed: 131 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
// under the License.
1717

1818
//! Regex expressions
19-
use arrow::array::{Array, ArrayRef, AsArray};
19+
use arrow::array::{Array, ArrayRef, AsArray, Datum};
2020
use arrow::compute::kernels::regexp;
2121
use arrow::datatypes::DataType;
2222
use arrow::datatypes::Field;
@@ -116,6 +116,14 @@ impl ScalarUDFImpl for RegexpMatchFunc {
116116

117117
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
118118
let args = &args.args;
119+
120+
// A literal pattern is the common case, and handing it to the kernel as
121+
// a scalar lets the regex be compiled once for the whole array. Any
122+
// other argument shape falls through to the general path below.
123+
if let Some(result) = regexp_match_scalar_pattern(args)? {
124+
return Ok(ColumnarValue::Array(result));
125+
}
126+
119127
let len = args
120128
.iter()
121129
.fold(Option::<usize>::None, |acc, arg| match arg {
@@ -145,6 +153,61 @@ impl ScalarUDFImpl for RegexpMatchFunc {
145153
}
146154
}
147155

156+
/// Runs `regexp_match` with the pattern (and flags, if given) passed to the
157+
/// kernel as scalar [`Datum`]s, so the regex is compiled once for the whole
158+
/// array.
159+
///
160+
/// Applies when the values are an array, the pattern is a non-null scalar of
161+
/// the same string type as the values, and the flags, if given, are a scalar of
162+
/// that same type and are not the unsupported "global" flag.
163+
///
164+
/// Returns `Ok(None)` for every other argument shape, leaving the caller's
165+
/// general path to materialize each argument as an array, zip the rows, and
166+
/// raise whatever error the shape warrants.
167+
fn regexp_match_scalar_pattern(args: &[ColumnarValue]) -> Result<Option<ArrayRef>> {
168+
let (values, pattern, flags) = match args {
169+
[values, pattern] => (values, pattern, None),
170+
[values, pattern, flags] => (values, pattern, Some(flags)),
171+
_ => return Ok(None),
172+
};
173+
174+
let (ColumnarValue::Array(values), ColumnarValue::Scalar(pattern)) =
175+
(values, pattern)
176+
else {
177+
return Ok(None);
178+
};
179+
let flags = match flags {
180+
// An array of flags has to be zipped with the values row by row.
181+
Some(ColumnarValue::Array(_)) => return Ok(None),
182+
Some(ColumnarValue::Scalar(flags)) => Some(flags),
183+
None => None,
184+
};
185+
186+
// The kernel requires the values, the pattern and the flags to share one
187+
// string type.
188+
let value_type = values.data_type();
189+
190+
if !matches!(pattern.try_as_str(), Some(Some(_)))
191+
|| &pattern.data_type() != value_type
192+
|| flags.is_some_and(|flags| {
193+
flags.try_as_str() == Some(Some("g")) || &flags.data_type() != value_type
194+
})
195+
{
196+
return Ok(None);
197+
}
198+
199+
let pattern = pattern.to_scalar()?;
200+
let flags = flags.map(ScalarValue::to_scalar).transpose()?;
201+
202+
regexp::regexp_match(
203+
values,
204+
&pattern,
205+
flags.as_ref().map(|flags| flags as &dyn Datum),
206+
)
207+
.map(Some)
208+
.map_err(|e| arrow_datafusion_err!(e))
209+
}
210+
148211
pub fn regexp_match(args: &[ArrayRef]) -> Result<ArrayRef> {
149212
match args.len() {
150213
2 => regexp::regexp_match(&args[0], &args[1], None)
@@ -257,4 +320,71 @@ mod tests {
257320
"Error during planning: regexp_match() does not support the \"global\" option"
258321
);
259322
}
323+
324+
/// The literal-pattern fast path must agree with the general path that
325+
/// zips a pattern array with the values, for every argument shape.
326+
#[test]
327+
fn test_scalar_pattern_matches_array_pattern() {
328+
use super::{RegexpMatchFunc, ScalarValue};
329+
use arrow::array::{Array, ArrayRef};
330+
use arrow::datatypes::{DataType, Field};
331+
use datafusion_common::config::ConfigOptions;
332+
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
333+
334+
let values = Arc::new(StringArray::from(vec![
335+
Some("abc"),
336+
Some("ABC"),
337+
None,
338+
Some(""),
339+
Some("a-b-c"),
340+
])) as ArrayRef;
341+
342+
for pattern in ["([a-z])(b)?", "^(A)", "no-match", "", "[a-z]+"] {
343+
for flags in [None, Some("i")] {
344+
let mut scalar_args = vec![
345+
ColumnarValue::Array(Arc::clone(&values)),
346+
ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern.to_string()))),
347+
];
348+
let mut array_args = vec![
349+
Arc::clone(&values),
350+
Arc::new(StringArray::from(vec![pattern; values.len()])) as ArrayRef,
351+
];
352+
if let Some(flags) = flags {
353+
scalar_args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
354+
flags.to_string(),
355+
))));
356+
array_args
357+
.push(Arc::new(StringArray::from(vec![flags; values.len()]))
358+
as ArrayRef);
359+
}
360+
361+
let arg_fields = scalar_args
362+
.iter()
363+
.enumerate()
364+
.map(|(idx, arg)| {
365+
Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
366+
})
367+
.collect();
368+
let actual = RegexpMatchFunc::new()
369+
.invoke_with_args(ScalarFunctionArgs {
370+
args: scalar_args,
371+
arg_fields,
372+
number_rows: values.len(),
373+
return_field: Field::new_list(
374+
"f",
375+
Field::new_list_field(DataType::Utf8, true),
376+
true,
377+
)
378+
.into(),
379+
config_options: Arc::new(ConfigOptions::default()),
380+
})
381+
.unwrap()
382+
.to_array(values.len())
383+
.unwrap();
384+
385+
let expected = regexp_match(&array_args).unwrap();
386+
assert_eq!(&actual, &expected, "pattern={pattern:?} flags={flags:?}");
387+
}
388+
}
389+
}
260390
}

0 commit comments

Comments
 (0)