Rich-T-kid commented on code in PR #26186:
URL: https://github.com/apache/datafusion/pull/26186#discussion_r4367840684
##########
datafusion/functions/src/regex/regexpreplace.rs:
##########
@@ -119,33 +112,89 @@ impl ScalarUDFImpl for RegexpReplaceFunc {
}
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- use DataType::*;
- Ok(match &arg_types[0] {
- LargeUtf8 | LargeBinary => LargeUtf8,
- Utf8 | Binary => Utf8,
- Utf8View | BinaryView => Utf8View,
- Null => Null,
- Dictionary(_, t) => match **t {
- LargeUtf8 | LargeBinary => LargeUtf8,
- Utf8 | Binary => Utf8,
- Null => Null,
- _ => {
- return plan_err!(
- "the regexp_replace can only accept strings but got
{:?}",
- **t
- );
- }
- },
- other => {
- return plan_err!(
- "The regexp_replace function can only accept strings. Got
{other}"
- );
+ transform_leaf_type_preserving_encoding(&arg_types[0],
®exp_replace_leaf_type)
+ }
+
+ fn return_field_from_args(&self, args: ReturnFieldArgs) ->
Result<FieldRef> {
+ let input = &args.arg_fields[0];
+ let others_scalar =
args.scalar_arguments.iter().skip(1).all(Option::is_some);
+ let output_type = match input.data_type() {
+ DataType::Dictionary(_, value_type) if !others_scalar => {
+ regexp_replace_leaf_type(value_type.as_ref())?
}
- })
+ other => self.return_type(std::slice::from_ref(other))?,
+ };
+ Ok(Arc::new(Field::new(self.name(), output_type, true)))
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+ if !(3..=4).contains(&arg_types.len()) {
+ return plan_err!(
+ "regexp_replace expects 3 or 4 arguments, got {}",
+ arg_types.len()
+ );
+ }
+ let leaf = match &arg_types[0] {
+ DataType::Dictionary(_, v) => v.as_ref(),
+ other => other,
+ };
+ let is_string_leaf = matches!(
+ leaf,
+ DataType::Utf8
+ | DataType::LargeUtf8
+ | DataType::Utf8View
+ | DataType::Binary
+ | DataType::LargeBinary
+ | DataType::BinaryView
+ );
+ let target = if is_string_leaf {
+ regexp_replace_leaf_type(leaf)?
+ } else {
+ DataType::Utf8View
+ };
+ let first = match &arg_types[0] {
+ DataType::Dictionary(k, _) if is_string_leaf => {
+ DataType::Dictionary(k.clone(), Box::new(target.clone()))
+ }
+ _ => target.clone(),
+ };
+ let mut out = Vec::with_capacity(arg_types.len());
+ out.push(first);
+ for _ in 1..arg_types.len() {
+ out.push(target.clone());
+ }
+ Ok(out)
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
- let args = &args.args;
+ let preserve_dictionary =
+ matches!(args.return_field.data_type(), DataType::Dictionary(_,
_));
+ let mut args = args.args;
+
+ let scalar_dict_key = if let Some(ColumnarValue::Scalar(
+ ScalarValue::Dictionary(key_type, value),
+ )) = args.first()
+ {
+ let key = key_type.clone();
+ args[0] = ColumnarValue::Scalar(value.as_ref().clone());
+ preserve_dictionary.then_some(key)
+ } else {
+ None
+ };
+
+ if let ColumnarValue::Array(array) = &args[0]
+ && let DataType::Dictionary(_, value_type) =
array.data_type().clone()
+ {
+ if preserve_dictionary {
+ let dict = array.as_any_dictionary();
Review Comment:
Done in 3f909972e. The dictionary is compacted when it has more values than
rows before regexp_replace runs.
##########
datafusion/sqllogictest/test_files/regexp/regexp_replace.slt:
##########
@@ -138,6 +138,41 @@ SELECT regexp_replace(arrow_cast('foobar',
'Dictionary(Int32, Utf8)'), 'bar', 'x
----
fooxx
+statement ok
+CREATE TABLE regexp_replace_dict AS
+SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS url
+FROM (VALUES
+ ('https://www.example.com/path'),
+ ('http://test.org/foo/bar'),
+ ('https://example.com/')
+) AS t(column1);
+
+query T
+SELECT DISTINCT arrow_typeof(regexp_replace(url,
'^https?://(?:www\.)?([^/]+)/.*$', '\1'))
+FROM regexp_replace_dict;
+----
+Dictionary(Int32, Utf8)
+
+statement ok
+DROP TABLE regexp_replace_dict;
+
+statement ok
+CREATE TABLE regexp_replace_dict_runtime_scalars AS
+SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS s,
+ column2 AS p
+FROM (VALUES ('abc', 'a'), ('abd', 'b')) AS v;
+
+query T
+SELECT regexp_replace(s, 'a', arrow_typeof(p))
+FROM regexp_replace_dict_runtime_scalars
+ORDER BY p;
+----
+Utf8bc
+Utf8bd
+
+statement ok
+DROP TABLE regexp_replace_dict_runtime_scalars;
Review Comment:
Added the LargeUtf8, Utf8View, Binary, and non-literal fallback cases in
3f909972e.
--
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]
