Rich-T-kid commented on code in PR #26186:
URL: https://github.com/apache/datafusion/pull/26186#discussion_r4308211956
##########
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:
thanks for catching this ,
https://github.com/apache/datafusion/pull/26186/commits/80bd16e2f47a5178eb97f49cbaf591f57fb91ff3
should address this
--
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]
