You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "olsemeno (via GitHub)" <gi...@apache.org> on 2026/03/17 09:55:26 UTC
[PR] Feat/regexp extract [datafusion]
olsemeno opened a new pull request, #20985:
URL: https://github.com/apache/datafusion/pull/20985
## Which issue does this PR close?
- Closes #.
## Rationale for this change
Apache DataFusion is missing a `regexp_extract` scalar function that is available in Apache Spark (via `functions.regexp_extract`). This function is commonly used in ETL pipelines to extract specific capture groups from strings using regular expressions. Adding it improves compatibility with Spark workloads migrating to DataFusion.
## What changes are included in this PR?
- Added a new scalar UDF `regexp_extract(str, regexp, idx)` in `datafusion/functions/src/regex/regexpextract.rs`
- Accepts `Utf8`, `LargeUtf8`, and `Utf8View` string types
- Third argument `idx` (`Int64`) selects the capture group: `0` returns the full match, `≥1` returns the corresponding group
- Returns an empty string `""` when the pattern does not match or the requested group is absent (Spark-compatible behavior)
- Returns `NULL` when the input string is `NULL`
- Returns an error for negative `idx` values
- Supports per-row `idx` values (column argument), not just scalar literals
- Uses the existing `compile_and_cache_regex` helper for regex compilation and caching
- Registered the function in `datafusion/functions/src/regex/mod.rs`
- Added `regexp_extract` usage examples to `datafusion-examples/examples/builtin_functions/regexp.rs`
## Are these changes tested?
Yes. Ten unit tests are included in `regexpextract.rs` covering:
| Test | Scenario |
|---|---|
| `test_basic_group` | Extract a specific capture group |
| `test_idx_zero_returns_whole_match` | `idx=0` returns the full match |
| `test_no_match_returns_empty_string` | No match → `""` |
| `test_null_input_returns_null` | NULL input → NULL output |
| `test_empty_pattern_returns_empty_string` | Empty pattern → `""` |
| `test_idx_out_of_range_returns_empty_string` | Group index beyond available groups → `""` |
| `test_negative_idx_returns_error` | Negative `idx` → error |
| `test_multiple_groups` | Multiple capture groups, select second |
| `test_idx_as_column` | `idx` as a per-row column (not scalar) |
| `test_batch_mixed_nulls` | Batch with mixed NULL and non-NULL rows |
## Are there any user-facing changes?
Yes. A new scalar function `regexp_extract(str, regexp, idx)` is now available in both SQL and the DataFrame API.
```sql
SELECT regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 1);
-- returns: '2024'
```
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feat/regexp extract [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #20985: Feat/regexp extract
URL: https://github.com/apache/datafusion/pull/20985
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feat/regexp extract [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #20985:
URL: https://github.com/apache/datafusion/pull/20985#issuecomment-4814722169
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feat/regexp extract [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #20985:
URL: https://github.com/apache/datafusion/pull/20985#issuecomment-4321347007
> I don't know of a reason those prior PRs weren't merged -- AFAICS this is a useful feature.
I think it was a combination of the PRs becoming inactive and some indecision on how to handle aligning with Spark semantics. I would be willing to help review a PR for this function, but I would prefer not having to review two PRs implementing the same functionality especially if there's a chance they again become inactive 🙁
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Feat/regexp extract [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on code in PR #20985:
URL: https://github.com/apache/datafusion/pull/20985#discussion_r3144661426
##########
datafusion/functions/src/regex/regexpextract.rs:
##########
@@ -0,0 +1,372 @@
+// 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 crate::regex::compile_and_cache_regex;
+use datafusion_macros::user_doc;
+use arrow::array::{Array, ArrayRef, AsArray, GenericByteViewArray, PrimitiveArray, StringBuilder, StringArrayType, StringViewBuilder};
+use arrow::datatypes::StringViewType;
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{Int64, LargeUtf8, Utf8, Utf8View};
+use arrow::datatypes::Int64Type;
+use datafusion_common::{Result, ScalarValue, arrow_datafusion_err, exec_err};
+use datafusion_expr::{
+ ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact, Volatility,
+};
+use std::any::Any;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+#[user_doc(
+ doc_section(label = "Regular Expression Functions"),
+ description = r#"Extracts the first match of a [regular expression](https://docs.rs/regex/latest/regex/#syntax) in a string, returning a specific capture group.
+If the regex did not match, or the specified group did not match, an empty string is returned."#,
+ syntax_example = "regexp_extract(str, regexp, idx)",
+ sql_example = r#"```sql
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(1)) |
++------------------------------------------------------------------+
+| 2024 |
++------------------------------------------------------------------+
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 0);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(0)) |
++------------------------------------------------------------------+
+| 2024-03-16 |
++------------------------------------------------------------------+
+> select regexp_extract('no digits here', '(\d+)', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("no digits here"),Utf8("(\d+)"),Int64(1)) |
++------------------------------------------------------------------+
+| |
++------------------------------------------------------------------+
+```
+Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/regexp.rs)
+"#,
+ standard_argument(name = "str", prefix = "String"),
+ argument(
+ name = "regexp",
+ description = "Regular expression to match against. Can be a constant, column, or function."
+ ),
+ argument(
+ name = "idx",
+ description = r#"The capture group index to return. Index **0** returns the entire matched string. Index **N** returns the Nth capture group `(...)`.
+If the group index exceeds the number of groups in the pattern, an empty string is returned."#
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct RegexpExtractFunc {
+ signature: Signature,
+}
+
+impl Default for RegexpExtractFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl RegexpExtractFunc {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::one_of(
+ vec![
+ Exact(vec![Utf8View, Utf8View, Int64]),
+ Exact(vec![Utf8, Utf8, Int64]),
+ Exact(vec![LargeUtf8, LargeUtf8, Int64]),
+ ],
+ Volatility::Immutable,
+ ),
+ }
+ }
+}
+
+impl ScalarUDFImpl for RegexpExtractFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "regexp_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ match &arg_types[0] {
+ DataType::Null => Ok(DataType::Null),
+ Utf8View => Ok(Utf8View),
+ LargeUtf8 => Ok(LargeUtf8),
+ _ => Ok(Utf8),
+ }
Review Comment:
```suggestion
Ok(arg_types[0].clone())
```
##########
datafusion/functions/src/regex/regexpextract.rs:
##########
@@ -0,0 +1,372 @@
+// 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 crate::regex::compile_and_cache_regex;
+use datafusion_macros::user_doc;
+use arrow::array::{Array, ArrayRef, AsArray, GenericByteViewArray, PrimitiveArray, StringBuilder, StringArrayType, StringViewBuilder};
+use arrow::datatypes::StringViewType;
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{Int64, LargeUtf8, Utf8, Utf8View};
+use arrow::datatypes::Int64Type;
+use datafusion_common::{Result, ScalarValue, arrow_datafusion_err, exec_err};
+use datafusion_expr::{
+ ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact, Volatility,
+};
+use std::any::Any;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+#[user_doc(
+ doc_section(label = "Regular Expression Functions"),
+ description = r#"Extracts the first match of a [regular expression](https://docs.rs/regex/latest/regex/#syntax) in a string, returning a specific capture group.
+If the regex did not match, or the specified group did not match, an empty string is returned."#,
+ syntax_example = "regexp_extract(str, regexp, idx)",
+ sql_example = r#"```sql
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(1)) |
++------------------------------------------------------------------+
+| 2024 |
++------------------------------------------------------------------+
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 0);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(0)) |
++------------------------------------------------------------------+
+| 2024-03-16 |
++------------------------------------------------------------------+
+> select regexp_extract('no digits here', '(\d+)', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("no digits here"),Utf8("(\d+)"),Int64(1)) |
++------------------------------------------------------------------+
+| |
++------------------------------------------------------------------+
+```
+Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/regexp.rs)
+"#,
+ standard_argument(name = "str", prefix = "String"),
+ argument(
+ name = "regexp",
+ description = "Regular expression to match against. Can be a constant, column, or function."
+ ),
+ argument(
+ name = "idx",
+ description = r#"The capture group index to return. Index **0** returns the entire matched string. Index **N** returns the Nth capture group `(...)`.
+If the group index exceeds the number of groups in the pattern, an empty string is returned."#
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct RegexpExtractFunc {
+ signature: Signature,
+}
+
+impl Default for RegexpExtractFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl RegexpExtractFunc {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::one_of(
+ vec![
+ Exact(vec![Utf8View, Utf8View, Int64]),
+ Exact(vec![Utf8, Utf8, Int64]),
+ Exact(vec![LargeUtf8, LargeUtf8, Int64]),
+ ],
+ Volatility::Immutable,
+ ),
+ }
+ }
+}
+
+impl ScalarUDFImpl for RegexpExtractFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "regexp_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ match &arg_types[0] {
+ DataType::Null => Ok(DataType::Null),
+ Utf8View => Ok(Utf8View),
+ LargeUtf8 => Ok(LargeUtf8),
+ _ => Ok(Utf8),
+ }
+ }
+
+ fn invoke_with_args(
+ &self,
+ args: datafusion_expr::ScalarFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ let args = &args.args;
+
+ let len = args
+ .iter()
+ .fold(Option::<usize>::None, |acc, arg| match arg {
+ ColumnarValue::Scalar(_) => acc,
+ ColumnarValue::Array(a) => Some(a.len()),
+ });
+
+ let is_scalar = len.is_none();
+ let inferred_length = len.unwrap_or(1);
+ let args = args
+ .iter()
+ .map(|arg| arg.to_array(inferred_length))
+ .collect::<Result<Vec<_>>>()?;
+
+ let result = regexp_extract_func(&args);
+ if is_scalar {
+ let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
+ result.map(ColumnarValue::Scalar)
+ } else {
+ result.map(ColumnarValue::Array)
+ }
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+pub fn regexp_extract_func(args: &[ArrayRef]) -> Result<ArrayRef> {
+ if args.len() != 3 {
+ return exec_err!(
+ "regexp_extract requires exactly 3 arguments, got {}",
+ args.len()
+ );
+ }
+
+ let idx_array = args[2].as_primitive::<Int64Type>();
+
+ match args[0].data_type() {
+ Utf8View => regexp_extract_inner_view(
+ args[0].as_string_view(),
+ args[1].as_string_view(),
+ idx_array,
+ ),
+ Utf8 => regexp_extract_inner(
+ &args[0].as_string::<i32>(),
+ &args[1].as_string::<i32>(),
+ idx_array,
+ ),
+ LargeUtf8 => regexp_extract_inner(
+ &args[0].as_string::<i64>(),
+ &args[1].as_string::<i64>(),
+ idx_array,
+ ),
+ other => exec_err!("regexp_extract does not support type {other:?}"),
+ }
+}
+
+
+fn regexp_extract_inner<'a, S>(
+ values: &S,
+ regex_array: &S,
+ idx_array: &PrimitiveArray<Int64Type>,
+) -> Result<ArrayRef>
+where
+ S: StringArrayType<'a>,
+{
+ let is_regex_scalar = regex_array.len() == 1;
+ let is_idx_scalar = idx_array.len() == 1;
+ let mut regex_cache = HashMap::new();
+ let mut builder = StringBuilder::with_capacity(values.len(), values.len() * 16);
+
+ for i in 0..values.len() {
+ let pattern_i = if is_regex_scalar { 0 } else { i };
+ let idx_i = if is_idx_scalar { 0 } else { i };
Review Comment:
Is this "is_scalar" logic strictly necessary?
##########
datafusion/functions/src/regex/regexpextract.rs:
##########
@@ -0,0 +1,372 @@
+// 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 crate::regex::compile_and_cache_regex;
+use datafusion_macros::user_doc;
+use arrow::array::{Array, ArrayRef, AsArray, GenericByteViewArray, PrimitiveArray, StringBuilder, StringArrayType, StringViewBuilder};
+use arrow::datatypes::StringViewType;
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{Int64, LargeUtf8, Utf8, Utf8View};
+use arrow::datatypes::Int64Type;
+use datafusion_common::{Result, ScalarValue, arrow_datafusion_err, exec_err};
+use datafusion_expr::{
+ ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact, Volatility,
+};
+use std::any::Any;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+#[user_doc(
+ doc_section(label = "Regular Expression Functions"),
+ description = r#"Extracts the first match of a [regular expression](https://docs.rs/regex/latest/regex/#syntax) in a string, returning a specific capture group.
+If the regex did not match, or the specified group did not match, an empty string is returned."#,
+ syntax_example = "regexp_extract(str, regexp, idx)",
+ sql_example = r#"```sql
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(1)) |
++------------------------------------------------------------------+
+| 2024 |
++------------------------------------------------------------------+
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 0);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(0)) |
++------------------------------------------------------------------+
+| 2024-03-16 |
++------------------------------------------------------------------+
+> select regexp_extract('no digits here', '(\d+)', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("no digits here"),Utf8("(\d+)"),Int64(1)) |
++------------------------------------------------------------------+
+| |
++------------------------------------------------------------------+
+```
+Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/regexp.rs)
+"#,
+ standard_argument(name = "str", prefix = "String"),
+ argument(
+ name = "regexp",
+ description = "Regular expression to match against. Can be a constant, column, or function."
+ ),
+ argument(
+ name = "idx",
+ description = r#"The capture group index to return. Index **0** returns the entire matched string. Index **N** returns the Nth capture group `(...)`.
+If the group index exceeds the number of groups in the pattern, an empty string is returned."#
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct RegexpExtractFunc {
+ signature: Signature,
+}
+
+impl Default for RegexpExtractFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl RegexpExtractFunc {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::one_of(
+ vec![
+ Exact(vec![Utf8View, Utf8View, Int64]),
+ Exact(vec![Utf8, Utf8, Int64]),
+ Exact(vec![LargeUtf8, LargeUtf8, Int64]),
+ ],
+ Volatility::Immutable,
+ ),
+ }
+ }
+}
+
+impl ScalarUDFImpl for RegexpExtractFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "regexp_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ match &arg_types[0] {
+ DataType::Null => Ok(DataType::Null),
+ Utf8View => Ok(Utf8View),
+ LargeUtf8 => Ok(LargeUtf8),
+ _ => Ok(Utf8),
+ }
+ }
+
+ fn invoke_with_args(
+ &self,
+ args: datafusion_expr::ScalarFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ let args = &args.args;
+
+ let len = args
+ .iter()
+ .fold(Option::<usize>::None, |acc, arg| match arg {
+ ColumnarValue::Scalar(_) => acc,
+ ColumnarValue::Array(a) => Some(a.len()),
+ });
+
+ let is_scalar = len.is_none();
+ let inferred_length = len.unwrap_or(1);
+ let args = args
+ .iter()
+ .map(|arg| arg.to_array(inferred_length))
+ .collect::<Result<Vec<_>>>()?;
+
+ let result = regexp_extract_func(&args);
+ if is_scalar {
+ let result = result.and_then(|arr| ScalarValue::try_from_array(&arr, 0));
+ result.map(ColumnarValue::Scalar)
+ } else {
+ result.map(ColumnarValue::Array)
+ }
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+pub fn regexp_extract_func(args: &[ArrayRef]) -> Result<ArrayRef> {
+ if args.len() != 3 {
+ return exec_err!(
+ "regexp_extract requires exactly 3 arguments, got {}",
+ args.len()
+ );
+ }
+
+ let idx_array = args[2].as_primitive::<Int64Type>();
+
+ match args[0].data_type() {
+ Utf8View => regexp_extract_inner_view(
+ args[0].as_string_view(),
+ args[1].as_string_view(),
+ idx_array,
+ ),
+ Utf8 => regexp_extract_inner(
+ &args[0].as_string::<i32>(),
+ &args[1].as_string::<i32>(),
+ idx_array,
+ ),
+ LargeUtf8 => regexp_extract_inner(
+ &args[0].as_string::<i64>(),
+ &args[1].as_string::<i64>(),
+ idx_array,
+ ),
+ other => exec_err!("regexp_extract does not support type {other:?}"),
+ }
+}
+
+
+fn regexp_extract_inner<'a, S>(
+ values: &S,
+ regex_array: &S,
+ idx_array: &PrimitiveArray<Int64Type>,
+) -> Result<ArrayRef>
+where
+ S: StringArrayType<'a>,
+{
+ let is_regex_scalar = regex_array.len() == 1;
+ let is_idx_scalar = idx_array.len() == 1;
+ let mut regex_cache = HashMap::new();
+ let mut builder = StringBuilder::with_capacity(values.len(), values.len() * 16);
+
+ for i in 0..values.len() {
+ let pattern_i = if is_regex_scalar { 0 } else { i };
+ let idx_i = if is_idx_scalar { 0 } else { i };
+
+ if values.is_null(i) || regex_array.is_null(pattern_i) || idx_array.is_null(idx_i) {
+ builder.append_null();
+ continue;
+ }
+
+ let extracted = process_row(
+ values.value(i),
+ regex_array.value(pattern_i),
+ idx_array.value(idx_i),
+ &mut regex_cache,
+ )?;
+ builder.append_value(extracted);
+ }
+
+ Ok(Arc::new(builder.finish()))
+}
+
+fn regexp_extract_inner_view<'a>(
+ values: &'a GenericByteViewArray<StringViewType>,
+ regex_array: &'a GenericByteViewArray<StringViewType>,
+ idx_array: &PrimitiveArray<Int64Type>,
+) -> Result<ArrayRef> {
+ let is_regex_scalar = regex_array.len() == 1;
+ let is_idx_scalar = idx_array.len() == 1;
+ let mut regex_cache = HashMap::new();
+ let mut builder = StringViewBuilder::with_capacity(values.len());
+
+ for i in 0..values.len() {
+ let pattern_i = if is_regex_scalar { 0 } else { i };
+ let idx_i = if is_idx_scalar { 0 } else { i };
+
+ if values.is_null(i) || regex_array.is_null(pattern_i) || idx_array.is_null(idx_i) {
+ builder.append_null();
+ continue;
+ }
+
+ let extracted = process_row(
+ values.value(i),
+ regex_array.value(pattern_i),
+ idx_array.value(idx_i),
+ &mut regex_cache,
+ )?;
+ builder.append_value(extracted);
+ }
+
+ Ok(Arc::new(builder.finish()))
+}
+
+fn process_row<'a>(
+ value: &'a str,
+ pattern: &'a str,
+ idx: i64,
+ regex_cache: &mut HashMap<(&'a str, Option<&'a str>), regex::Regex>,
+) -> Result<&'a str> {
+ if idx < 0 {
+ return exec_err!("regexp_extract requires a non-negative idx, got {idx}");
+ }
+ if pattern.is_empty() {
+ return Ok("");
+ }
+ let re = compile_and_cache_regex(pattern, None, regex_cache)
+ .map_err(|e| arrow_datafusion_err!(e))?;
+ Ok(extract_group(value, re, idx as usize))
+}
+
+fn extract_group<'a>(value: &'a str, re: ®ex::Regex, idx: usize) -> &'a str {
+ match re.captures(value) {
+ None => "",
+ Some(caps) => caps.get(idx).map_or("", |m| m.as_str()),
+ }
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
Could we move these tests to be in SLTs?
##########
datafusion/functions/src/regex/regexpextract.rs:
##########
@@ -0,0 +1,372 @@
+// 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 crate::regex::compile_and_cache_regex;
+use datafusion_macros::user_doc;
+use arrow::array::{Array, ArrayRef, AsArray, GenericByteViewArray, PrimitiveArray, StringBuilder, StringArrayType, StringViewBuilder};
+use arrow::datatypes::StringViewType;
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{Int64, LargeUtf8, Utf8, Utf8View};
+use arrow::datatypes::Int64Type;
+use datafusion_common::{Result, ScalarValue, arrow_datafusion_err, exec_err};
+use datafusion_expr::{
+ ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature::Exact, Volatility,
+};
+use std::any::Any;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+#[user_doc(
+ doc_section(label = "Regular Expression Functions"),
+ description = r#"Extracts the first match of a [regular expression](https://docs.rs/regex/latest/regex/#syntax) in a string, returning a specific capture group.
+If the regex did not match, or the specified group did not match, an empty string is returned."#,
+ syntax_example = "regexp_extract(str, regexp, idx)",
+ sql_example = r#"```sql
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(1)) |
++------------------------------------------------------------------+
+| 2024 |
++------------------------------------------------------------------+
+> select regexp_extract('2024-03-16', '(\d{4})-(\d{2})-(\d{2})', 0);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("2024-03-16"),Utf8("(\d{4})-(\d{2})-(\d{2})"),Int64(0)) |
++------------------------------------------------------------------+
+| 2024-03-16 |
++------------------------------------------------------------------+
+> select regexp_extract('no digits here', '(\d+)', 1);
++------------------------------------------------------------------+
+| regexp_extract(Utf8("no digits here"),Utf8("(\d+)"),Int64(1)) |
++------------------------------------------------------------------+
+| |
++------------------------------------------------------------------+
+```
+Additional examples can be found [here](https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/builtin_functions/regexp.rs)
+"#,
+ standard_argument(name = "str", prefix = "String"),
+ argument(
+ name = "regexp",
+ description = "Regular expression to match against. Can be a constant, column, or function."
+ ),
+ argument(
+ name = "idx",
+ description = r#"The capture group index to return. Index **0** returns the entire matched string. Index **N** returns the Nth capture group `(...)`.
+If the group index exceeds the number of groups in the pattern, an empty string is returned."#
+ )
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct RegexpExtractFunc {
+ signature: Signature,
+}
+
+impl Default for RegexpExtractFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl RegexpExtractFunc {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::one_of(
+ vec![
+ Exact(vec![Utf8View, Utf8View, Int64]),
+ Exact(vec![Utf8, Utf8, Int64]),
+ Exact(vec![LargeUtf8, LargeUtf8, Int64]),
+ ],
+ Volatility::Immutable,
+ ),
+ }
+ }
+}
+
+impl ScalarUDFImpl for RegexpExtractFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "regexp_extract"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ match &arg_types[0] {
+ DataType::Null => Ok(DataType::Null),
+ Utf8View => Ok(Utf8View),
+ LargeUtf8 => Ok(LargeUtf8),
+ _ => Ok(Utf8),
+ }
+ }
+
+ fn invoke_with_args(
+ &self,
+ args: datafusion_expr::ScalarFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ let args = &args.args;
+
+ let len = args
Review Comment:
[`make_scalar_function`](https://docs.rs/datafusion/latest/datafusion/functions/utils/fn.make_scalar_function.html) already does all this logic for us, so should reuse that here
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org