Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

improve performance of regexp_count #13364

Merged
merged 3 commits into from
Nov 12, 2024
Merged

improve performance of regexp_count #13364

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5d93da0
improve performance of regexp_count
Nov 12, 2024
fba93c3
fix clippy
Nov 12, 2024
3bd10d1
collect with Int64Array to eliminate one temp Vec
Nov 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 39 additions & 43 deletions datafusion/functions/src/regex/regexpcount.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ use datafusion_expr::{
};
use itertools::izip;
use regex::Regex;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::sync::{Arc, OnceLock};

Expand Down Expand Up @@ -312,12 +311,12 @@ where

let pattern = compile_regex(regex, flags_scalar)?;

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
values
.iter()
.map(|value| count_matches(value, &pattern, start_scalar))
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(true, true, false) => {
let regex = match regex_scalar {
Expand All @@ -336,17 +335,17 @@ where
)));
}

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
values
.iter()
.zip(flags_array.iter())
.map(|(value, flags)| {
let pattern =
compile_and_cache_regex(regex, flags, &mut regex_cache)?;
count_matches(value, &pattern, start_scalar)
count_matches(value, pattern, start_scalar)
Dandandan marked this conversation as resolved.
Show resolved Hide resolved
})
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(true, false, true) => {
let regex = match regex_scalar {
Expand All @@ -360,13 +359,13 @@ where

let start_array = start_array.unwrap();

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
values
.iter()
.zip(start_array.iter())
.map(|(value, start)| count_matches(value, &pattern, start))
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(true, false, false) => {
let regex = match regex_scalar {
Expand All @@ -385,7 +384,7 @@ where
)));
}

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
izip!(
values.iter(),
start_array.unwrap().iter(),
Expand All @@ -395,10 +394,10 @@ where
let pattern =
compile_and_cache_regex(regex, flags, &mut regex_cache)?;

count_matches(value, &pattern, start)
count_matches(value, pattern, start)
Dandandan marked this conversation as resolved.
Show resolved Hide resolved
})
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(false, true, true) => {
if values.len() != regex_array.len() {
Expand All @@ -409,7 +408,7 @@ where
)));
}

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
values
.iter()
.zip(regex_array.iter())
Expand All @@ -424,10 +423,10 @@ where
flags_scalar,
&mut regex_cache,
)?;
count_matches(value, &pattern, start_scalar)
count_matches(value, pattern, start_scalar)
})
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(false, true, false) => {
if values.len() != regex_array.len() {
Expand All @@ -447,7 +446,7 @@ where
)));
}

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
izip!(values.iter(), regex_array.iter(), flags_array.iter())
.map(|(value, regex, flags)| {
let regex = match regex {
Expand All @@ -458,10 +457,10 @@ where
let pattern =
compile_and_cache_regex(regex, flags, &mut regex_cache)?;

count_matches(value, &pattern, start_scalar)
count_matches(value, pattern, start_scalar)
})
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(false, false, true) => {
if values.len() != regex_array.len() {
Expand All @@ -481,7 +480,7 @@ where
)));
}

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
izip!(values.iter(), regex_array.iter(), start_array.iter())
.map(|(value, regex, start)| {
let regex = match regex {
Expand All @@ -494,10 +493,10 @@ where
flags_scalar,
&mut regex_cache,
)?;
count_matches(value, &pattern, start)
count_matches(value, pattern, start)
})
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
(false, false, false) => {
if values.len() != regex_array.len() {
Expand Down Expand Up @@ -526,7 +525,7 @@ where
)));
}

Ok(Arc::new(Int64Array::from_iter_values(
Ok(Arc::new(
izip!(
values.iter(),
regex_array.iter(),
Expand All @@ -541,27 +540,24 @@ where

let pattern =
compile_and_cache_regex(regex, flags, &mut regex_cache)?;
count_matches(value, &pattern, start)
count_matches(value, pattern, start)
})
.collect::<Result<Vec<i64>, ArrowError>>()?,
)))
.collect::<Result<Int64Array, ArrowError>>()?,
))
}
}
}

fn compile_and_cache_regex(
regex: &str,
flags: Option<&str>,
regex_cache: &mut HashMap<String, Regex>,
) -> Result<Regex, ArrowError> {
match regex_cache.entry(regex.to_string()) {
Entry::Vacant(entry) => {
let compiled = compile_regex(regex, flags)?;
entry.insert(compiled.clone());
Ok(compiled)
}
Entry::Occupied(entry) => Ok(entry.get().to_owned()),
fn compile_and_cache_regex<'a>(
regex: &'a str,
flags: Option<&'a str>,
regex_cache: &'a mut HashMap<String, Regex>,
Dimchikkk marked this conversation as resolved.
Show resolved Hide resolved
) -> Result<&'a Regex, ArrowError> {
if !regex_cache.contains_key(regex) {
let compiled = compile_regex(regex, flags)?;
regex_cache.insert(regex.to_string(), compiled);
}
Ok(regex_cache.get(regex).unwrap())
}

fn compile_regex(regex: &str, flags: Option<&str>) -> Result<Regex, ArrowError> {
Expand Down