jorgecarleitao commented on a change in pull request #9376:
URL: https://github.com/apache/arrow/pull/9376#discussion_r568343902
##########
File path: rust/datafusion/examples/simple_udf.rs
##########
@@ -54,50 +58,76 @@ fn create_context() -> Result<ExecutionContext> {
Ok(ctx)
}
+// a small utility function to compute pow(base, exponent)
+fn maybe_pow(base: &Option<f64>, exponent: &Option<f64>) -> Option<f64> {
+ match (base, exponent) {
+ // in arrow, any value can be null.
+ // Here we decide to make our UDF to return null when either base or
exponent is null.
+ (Some(base), Some(exponent)) => Some(base.powf(*exponent)),
+ _ => None,
+ }
+}
+
+fn pow_array(base: &dyn Array, exponent: &dyn Array) -> Result<ArrayRef> {
+ // 1. cast both arguments to f64. These casts MUST be aligned with the
signature or this function panics!
+ let base = base
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .expect("cast failed");
+ let exponent = exponent
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .expect("cast failed");
+
+ // this is guaranteed by DataFusion. We place it just to make it obvious.
+ assert_eq!(exponent.len(), base.len());
+
+ // 2. perform the computation
+ let array = base
+ .iter()
+ .zip(exponent.iter())
+ .map(|(base, exponent)| maybe_pow(&base, &exponent))
+ .collect::<Float64Array>();
+
+ // `Ok` because no error occurred during the calculation (we should add
one if exponent was [0, 1[ and the base < 0 because that panics!)
Review comment:
This represents an open interval: `0 <= exponent < 1` because `(x)^1 =
x` for all x (including negative ones).
----------------------------------------------------------------
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]
