diff --git a/datafusion/functions/src/regex/regexpreplace.rs b/datafusion/functions/src/regex/regexpreplace.rs index 176f704a65f35..40c74f0a6ab42 100644 --- a/datafusion/functions/src/regex/regexpreplace.rs +++ b/datafusion/functions/src/regex/regexpreplace.rs @@ -16,6 +16,8 @@ // under the License. //! Regex expressions +use memchr::memchr; + use arrow::array::ArrayDataBuilder; use arrow::array::BufferBuilder; use arrow::array::GenericStringArray; @@ -40,7 +42,8 @@ use datafusion_expr::{ Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; use datafusion_macros::user_doc; -use regex::Regex; +use regex::{CaptureLocations, Regex}; +use std::borrow::Cow; use std::collections::HashMap; use std::sync::{Arc, LazyLock}; @@ -199,6 +202,96 @@ fn regex_replace_posix_groups(replacement: &str) -> String { .into_owned() } +struct ShortRegex { + /// Shortened anchored regex used to extract capture group 1 directly. + /// See [`try_build_short_extract_regex`] for details. + short_re: Regex, + /// Reusable capture locations for `short_re` to avoid per-row allocation. + locs: CaptureLocations, +} + +/// Holds the normal compiled regex together with the optional fast path used +/// for `regexp_replace(str, '^...(capture)...*$', '\1')`. +struct OptimizedRegex { + /// Full regex used for the normal replacement path and as a correctness fallback. + re: Regex, + /// Precomputed state for the direct-extraction fast path, when applicable. + short_re: Option, +} + +impl OptimizedRegex { + /// Builds any reusable state needed by the extraction fast path. + /// + /// The fast path is only enabled for single replacements where the pattern + /// and replacement satisfy [`try_build_short_extract_regex`]. + fn new(re: Regex, limit: usize, pattern: &str, replacement: &str) -> Self { + let short_re = if limit == 1 { + try_build_short_extract_regex(pattern, replacement) + } else { + None + }; + + let short_re = short_re.map(|short_re| { + let locs = short_re.capture_locations(); + ShortRegex { short_re, locs } + }); + + Self { re, short_re } + } + + /// Applies the direct-extraction fast path when it preserves the result of + /// `Regex::replacen`; otherwise falls back to the full regex replacement. + fn replacen<'a>( + &mut self, + val: &'a str, + limit: usize, + replacement: &str, + ) -> Cow<'a, str> { + // If this pattern is not eligible for direct extraction, use the full regex. + let Some(ShortRegex { short_re, locs }) = self.short_re.as_mut() else { + return self.re.replacen(val, limit, replacement); + }; + + // If the shortened regex does not match, the original anchored regex would + // also leave the input unchanged. + if short_re.captures_read(locs, val).is_none() { + return Cow::Borrowed(val); + }; + + // `captures_read` succeeded, so the overall shortened match is present. + let match_end = locs.get(0).unwrap().1; + if memchr(b'\n', &val.as_bytes()[match_end..]).is_some() { + // If there is a newline after the match, we can't use the short + // regex since it won't match across lines. Fall back to the full + // regex replacement. + return self.re.replacen(val, limit, replacement); + }; + // The fast path only applies to `${1}` replacements, so the result is + // either capture group 1 or the empty string if that group did not match. + if let Some((start, end)) = locs.get(1) { + Cow::Borrowed(&val[start..end]) + } else { + Cow::Borrowed("") + } + } +} + +/// For anchored patterns like `^...(capture)....*$` where the replacement +/// is `\1`, build a shorter regex (stripping trailing `.*$`) and use +/// `captures_read` with `CaptureLocations` for direct extraction — no +/// `expand()`, no `String` allocation. +fn try_build_short_extract_regex(pattern: &str, replacement: &str) -> Option { + if replacement != "${1}" || !pattern.starts_with('^') || !pattern.ends_with(".*$") { + return None; + } + let short = &pattern[..pattern.len() - 3]; + let re = Regex::new(short).ok()?; + if re.captures_len() != 2 { + return None; + } + Some(re) +} + /// Replaces substring(s) matching a PCRE-like regular expression. /// /// The full list of supported features and syntax can be found at @@ -422,7 +515,7 @@ macro_rules! fetch_string_arg { /// hold a single Regex object for the replace operation. This also speeds /// up the pre-processing time of the replacement string, since it only /// needs to processed once. -fn _regexp_replace_static_pattern_replace( +fn regexp_replace_static_pattern_replace( args: &[ArrayRef], ) -> Result { let array_size = args[0].len(); @@ -457,6 +550,8 @@ fn _regexp_replace_static_pattern_replace( // with rust ones. let replacement = regex_replace_posix_groups(replacement); + let mut opt_re = OptimizedRegex::new(re, limit, &pattern, &replacement); + let string_array_type = args[0].data_type(); match string_array_type { DataType::Utf8 | DataType::LargeUtf8 => { @@ -475,7 +570,7 @@ fn _regexp_replace_static_pattern_replace( string_array.iter().for_each(|val| { if let Some(val) = val { - let result = re.replacen(val, limit, replacement.as_str()); + let result = opt_re.replacen(val, limit, replacement.as_str()); vals.append_slice(result.as_bytes()); } new_offsets.append(T::from_usize(vals.len()).unwrap()); @@ -496,8 +591,8 @@ fn _regexp_replace_static_pattern_replace( for val in string_view_array.iter() { if let Some(val) = val { - let result = re.replacen(val, limit, replacement.as_str()); - builder.append_value(result); + let result = opt_re.replacen(val, limit, replacement.as_str()); + builder.append_value(result.as_ref()); } else { builder.append_null(); } @@ -576,7 +671,7 @@ fn specialize_regexp_replace( arg.to_array(expansion_len) }) .collect::>>()?; - _regexp_replace_static_pattern_replace::(&args) + regexp_replace_static_pattern_replace::(&args) } // If there are no specialized implementations, we'll fall back to the @@ -710,7 +805,7 @@ mod tests { let replacements = <$T>::from(replacement); let expected = <$T>::from(expected); - let re = _regexp_replace_static_pattern_replace::<$O>(&[ + let re = regexp_replace_static_pattern_replace::<$O>(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -755,7 +850,7 @@ mod tests { let flags = StringArray::from(vec!["i"; 5]); let expected = <$T>::from(expected); - let re = _regexp_replace_static_pattern_replace::<$O>(&[ + let re = regexp_replace_static_pattern_replace::<$O>(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -787,7 +882,7 @@ mod tests { let replacements = StringArray::from(vec!["foo"; 5]); let expected = StringArray::from(vec![None::<&str>; 5]); - let re = _regexp_replace_static_pattern_replace::(&[ + let re = regexp_replace_static_pattern_replace::(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -804,7 +899,7 @@ mod tests { let replacements = StringArray::from(Vec::>::new()); let expected = StringArray::from(Vec::>::new()); - let re = _regexp_replace_static_pattern_replace::(&[ + let re = regexp_replace_static_pattern_replace::(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -822,7 +917,7 @@ mod tests { let flags = StringArray::from(vec![None::<&str>; 5]); let expected = StringArray::from(vec![None::<&str>; 5]); - let re = _regexp_replace_static_pattern_replace::(&[ + let re = regexp_replace_static_pattern_replace::(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -841,7 +936,7 @@ mod tests { let patterns = StringArray::from(vec!["["; 5]); let replacements = StringArray::from(vec!["foo"; 5]); - let re = _regexp_replace_static_pattern_replace::(&[ + let re = regexp_replace_static_pattern_replace::(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -878,7 +973,7 @@ mod tests { Some("c"), ]); - let re = _regexp_replace_static_pattern_replace::(&[ + let re = regexp_replace_static_pattern_replace::(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), @@ -906,7 +1001,7 @@ mod tests { let replacements = StringArray::from(vec!["foo"; 1]); let expected = StringArray::from(vec![Some("b"), None, Some("foo"), None, None]); - let re = _regexp_replace_static_pattern_replace::(&[ + let re = regexp_replace_static_pattern_replace::(&[ Arc::new(values), Arc::new(patterns), Arc::new(replacements), diff --git a/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt b/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt index a2eccfce5f695..99e0b10430186 100644 --- a/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt +++ b/datafusion/sqllogictest/test_files/regexp/regexp_replace.slt @@ -128,6 +128,35 @@ from (values ('a'), ('b')) as tbl(col); NULL NULL NULL NULL NULL NULL +# Extract domain from URL using anchored pattern with trailing .* +# This tests that the full URL suffix is replaced, not just the matched prefix +query T +SELECT regexp_replace(url, '^https?://(?:www\.)?([^/]+)/.*$', '\1') FROM (VALUES + ('https://www.example.com/path/to/page?q=1'), + ('http://test.org/foo/bar'), + ('https://example.com/'), + ('not-a-url') +) AS t(url); +---- +example.com +test.org +example.com +not-a-url + +# More than one capture group should disable the short-regex fast path. +# This still uses replacement \1, but captures_len() will be > 2, so the +# implementation must fall back to the normal regexp_replace path. +query T +SELECT regexp_replace(url, '^https?://((www\.)?([^/]+))/.*$', '\1') FROM (VALUES + ('https://www.example.com/path/to/page?q=1'), + ('http://test.org/foo/bar'), + ('not-a-url') +) AS t(url); +---- +www.example.com +test.org +not-a-url + # If the overall pattern matches but capture group 1 does not participate, # regexp_replace(..., '\1') should substitute the empty string, not keep # the original input.