andygrove commented on code in PR #23540:
URL: https://github.com/apache/datafusion/pull/23540#discussion_r3604574646
##########
datafusion/functions/src/regex/regexpinstr.rs:
##########
@@ -286,120 +286,104 @@ fn regexp_instr_inner<'a, S>(
regex_array: &S,
start_array: Option<&Int64Array>,
nth_array: Option<&Int64Array>,
- flags_array: Option<S>,
+ flags_array: Option<&S>,
subexp_array: Option<&Int64Array>,
) -> Result<ArrayRef, ArrowError>
where
S: StringArrayType<'a>,
{
let len = values.len();
+ let mut regex_cache = RegexCache::default();
+ let mut result = Int64Builder::with_capacity(len);
- let default_start_array = PrimitiveArray::<Int64Type>::from(vec![1; len]);
- let start_array = start_array.unwrap_or(&default_start_array);
- let start_input: Vec<i64> = (0..start_array.len())
- .map(|i| start_array.value(i)) // handle nulls as 0
- .collect();
-
- let default_nth_array = PrimitiveArray::<Int64Type>::from(vec![1; len]);
- let nth_array = nth_array.unwrap_or(&default_nth_array);
- let nth_input: Vec<i64> = (0..nth_array.len())
- .map(|i| nth_array.value(i)) // handle nulls as 0
- .collect();
-
- let flags_input = match flags_array {
- Some(flags) => flags.iter().collect(),
- None => vec![None; len],
- };
+ for i in 0..len {
+ if regex_array.is_null(i) {
+ result.append_null();
+ continue;
+ }
+ let regex = regex_array.value(i);
+ if regex.is_empty() {
+ result.append_value(0);
+ continue;
+ }
- let default_subexp_array = PrimitiveArray::<Int64Type>::from(vec![0; len]);
- let subexp_array = subexp_array.unwrap_or(&default_subexp_array);
- let subexp_input: Vec<i64> = (0..subexp_array.len())
- .map(|i| subexp_array.value(i)) // handle nulls as 0
- .collect();
-
- let mut regex_cache = HashMap::new();
-
- let result: Result<Vec<Option<i64>>, ArrowError> = izip!(
- values.iter(),
- regex_array.iter(),
- start_input.iter(),
- nth_input.iter(),
- flags_input.iter(),
- subexp_input.iter()
- )
- .map(|(value, regex, start, nth, flags, subexp)| match regex {
- None => Ok(None),
- Some("") => Ok(Some(0)),
- Some(regex) => get_index(
- value,
- regex,
- *start,
- *nth,
- *subexp,
- *flags,
- &mut regex_cache,
- ),
- })
- .collect();
- Ok(Arc::new(Int64Array::from(result?)))
-}
+ if values.is_null(i) {
+ result.append_null();
+ continue;
+ }
+ let value = values.value(i);
+ if value.is_empty() {
+ result.append_value(0);
+ continue;
+ }
-fn handle_subexp(
- pattern: &Regex,
- search_slice: &str,
- subexpr: i64,
- value: &str,
- byte_start_offset: usize,
-) -> Result<Option<i64>, ArrowError> {
- if let Some(captures) = pattern.captures(search_slice)
- && let Some(matched) = captures.get(subexpr as usize)
- {
- // Convert byte offset relative to search_slice back to 1-based
character offset
- // relative to the original `value` string.
- let start_char_offset =
- value[..byte_start_offset + matched.start()].chars().count() as
i64 + 1;
- return Ok(Some(start_char_offset));
+ let flags = match flags_array {
+ Some(flags) if !flags.is_null(i) => Some(flags.value(i)),
+ _ => None,
+ };
+ let pattern = regex_cache.get_or_compile(regex, flags)?;
+
+ // The defaults apply when the optional argument was not supplied at
+ // all. A supplied but null slot reads through as its raw buffer value.
+ let start = start_array.map_or(1, |array| array.value(i));
+ let nth = nth_array.map_or(1, |array| array.value(i));
+ let subexp = subexp_array.map_or(0, |array| array.value(i));
+
+ result.append_value(get_index(value, pattern, start, nth, subexp)?);
}
- Ok(Some(0)) // Return 0 if the subexpression was not found
+
+ Ok(Arc::new(result.finish()))
}
-fn get_nth_match(
- pattern: &Regex,
- search_slice: &str,
- n: i64,
- byte_start_offset: usize,
- value: &str,
-) -> Result<Option<i64>, ArrowError> {
- if let Some(mat) = pattern.find_iter(search_slice).nth((n - 1) as usize) {
- // Convert byte offset relative to search_slice back to 1-based
character offset
- // relative to the original `value` string.
- let match_start_byte_offset = byte_start_offset + mat.start();
- let match_start_char_offset =
- value[..match_start_byte_offset].chars().count() as i64 + 1;
- Ok(Some(match_start_char_offset))
- } else {
- Ok(Some(0)) // Return 0 if the N-th match was not found
+/// Compiles the patterns seen so far, keyed by `(pattern, flags)`.
+///
+/// Patterns are addressed by index rather than by reference so that `last` can
+/// memoize the previous row's pattern without holding a borrow of `indices`
+/// across rows. A literal pattern yields the same string on every row, so that
+/// memo means the common case never hashes a key.
+#[derive(Default)]
+struct RegexCache<'a> {
+ compiled: Vec<Regex>,
+ indices: HashMap<(&'a str, Option<&'a str>), usize>,
+ last: Option<((&'a str, Option<&'a str>), usize)>,
+}
+
+impl<'a> RegexCache<'a> {
+ fn get_or_compile(
+ &mut self,
+ regex: &'a str,
+ flags: Option<&'a str>,
+ ) -> Result<&Regex, ArrowError> {
+ let key = (regex, flags);
+ let index = match self.last {
+ Some((last_key, index)) if last_key == key => index,
+ _ => {
+ let index = match self.indices.entry(key) {
+ Entry::Occupied(entry) => *entry.get(),
+ Entry::Vacant(entry) => {
+ self.compiled.push(compile_regex(regex, flags)?);
+ *entry.insert(self.compiled.len() - 1)
+ }
+ };
+ self.last = Some((key, index));
+ index
+ }
+ };
+ Ok(&self.compiled[index])
}
}
-fn get_index<'strings, 'cache>(
- value: Option<&str>,
- pattern: &'strings str,
+
+/// Returns the 1-based character position of the `n`-th match of `pattern` in
+/// `value`, or 0 if there is no such match. The search begins at the 1-based
+/// character position `start`. A positive `subexpr` selects that capture group
+/// of the first match instead of the `n`-th match. `value` is non-empty.
+fn get_index(
+ value: &str,
+ pattern: &Regex,
start: i64,
n: i64,
subexpr: i64,
- flags: Option<&'strings str>,
- regex_cache: &'cache mut HashMap<(&'strings str, Option<&'strings str>),
Regex>,
-) -> Result<Option<i64>, ArrowError>
-where
- 'strings: 'cache,
-{
- let value = match value {
- None => return Ok(None),
- Some("") => return Ok(Some(0)),
- Some(value) => value,
- };
- let pattern: &Regex = compile_and_cache_regex(pattern, flags,
regex_cache)?;
- // println!("get_index: value = {}, pattern = {}, start = {}, n = {},
subexpr = {}, flags = {:?}", value, pattern, start, n, subexpr, flags);
+) -> Result<i64, ArrowError> {
Review Comment:
This PR doesn't make any functional changes. The index is 1-based, I believe.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
