Dandandan commented on a change in pull request #9376:
URL: https://github.com/apache/arrow/pull/9376#discussion_r571419011
##########
File path: rust/datafusion/src/physical_plan/functions.rs
##########
@@ -622,21 +560,66 @@ impl PhysicalExpr for ScalarFunctionExpr {
let inputs = self
.args
.iter()
- .map(|e| e.evaluate(batch).map(|v| v.into_array(batch.num_rows())))
+ .map(|e| e.evaluate(batch))
.collect::<Result<Vec<_>>>()?;
// evaluate the function
let fun = self.fun.as_ref();
- (fun)(&inputs).map(|a| ColumnarValue::Array(a))
+ (fun)(&inputs)
}
}
+/// decorates a function to handle [`ScalarValue`]s by coverting them to
arrays before calling the function
+/// and vice-versa after evaluation.
+pub fn make_scalar_function<F>(inner: F) -> ScalarFunctionImplementation
+where
+ F: Fn(&[ArrayRef]) -> Result<ArrayRef> + Sync + Send + 'static,
+{
+ Arc::new(move |args: &[ColumnarValue]| {
+ // first, identify if any of the arguments is an Array. If yes, store
its `len`,
+ // as any scalar will need to be converted to an array of len `len`.
+ let len = args
+ .iter()
+ .fold(Option::<usize>::None, |acc, arg| match arg {
+ ColumnarValue::Scalar(_) => acc,
+ ColumnarValue::Array(a) => Some(a.len()),
+ });
+
+ // to array
+ let args = if let Some(len) = len {
+ args.iter()
+ .map(|arg| arg.clone().into_array(len))
+ .collect::<Vec<ArrayRef>>()
+ } else {
+ args.iter()
+ .map(|arg| arg.clone().into_array(1))
+ .collect::<Vec<ArrayRef>>()
+ };
+
+ let result = (inner)(&args);
Review comment:
```suggestion
let result = inner(&args);
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
