kumarUjjawal commented on code in PR #26186:
URL: https://github.com/apache/datafusion/pull/26186#discussion_r4306691893
##########
datafusion/functions/src/regex/regexpreplace.rs:
##########
@@ -119,33 +112,90 @@ 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 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());
+ Some(key)
+ } else {
+ None
+ };
+
+ if let ColumnarValue::Array(array) = &args[0]
+ && let DataType::Dictionary(_, value_type) =
array.data_type().clone()
+ {
+ let others_scalar = args[1..]
+ .iter()
+ .all(|a| matches!(a, ColumnarValue::Scalar(_)));
+ if others_scalar {
+ let dict = array.as_any_dictionary();
+ let mut inner: Vec<ColumnarValue> =
Vec::with_capacity(args.len());
+ inner.push(ColumnarValue::Array(Arc::clone(dict.values())));
+ inner.extend_from_slice(&args[1..]);
+ let values_out = regexp_replace_func(&inner)?;
+ return Ok(ColumnarValue::Array(dict.with_values(values_out)));
Review Comment:
Could we use `args.return_field` to decide whether to keep dictionary
encoding? A nonliteral expression can return a scalar at runtime, even when
planning expects `Utf8`.
This SQL reproduces the failure:
```sql
CREATE TABLE t AS
SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS s,
column2 AS p
FROM (VALUES ('abc', 'a'), ('abd', 'b')) AS v;
SELECT regexp_replace(s, 'a', arrow_typeof(p)) FROM t;
```
The pre-PR function returns `Utf8bc` and `Utf8bd`. This version fails:
```text
Function 'regexp_replace' returned value of type 'Dictionary(Int32, Utf8)'
while the following type was promised at planning time and expected: 'Utf8'.
```
Please decode the result when the planned return type is a plain string, and
add this as a regression test.
--
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]
