alamb commented on a change in pull request #1510:
URL: https://github.com/apache/arrow-datafusion/pull/1510#discussion_r777007489
##########
File path: datafusion/src/physical_plan/math_expressions.rs
##########
@@ -102,6 +103,31 @@ math_unary_function!("ln", ln);
math_unary_function!("log2", log2);
math_unary_function!("log10", log10);
+/// factorial SQL function
+pub fn factorial(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ match &args[0] {
+ ColumnarValue::Array(array) => {
+ let x1 = array.as_any().downcast_ref::<Float64Array>();
+ match x1 {
+ Some(array) => {
+ let res: Float64Array =
+ arrow::compute::kernels::arity::unary(array, |x| {
+ factorial::factorial(x as u64)
+ });
+ let arc1 = Arc::new(res);
+ Ok(ColumnarValue::Array(arc1))
+ }
+ _ => Err(DataFusionError::Internal(
+ "Invalid data type for factorial function".to_string(),
+ )),
+ }
+ }
+ _ => Err(DataFusionError::Internal(
Review comment:
`ColumnValue::Scalar` is legit too I think (aka called with a constant)
##########
File path: datafusion/tests/sql/functions.rs
##########
@@ -17,6 +17,41 @@
use super::*;
+#[tokio::test]
+async fn factorial() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Float64,
true)]));
+
+ let data = RecordBatch::try_new(
+ schema.clone(),
+ vec![Arc::new(Float64Array::from(vec![
+ Some(4.0),
+ Some(0.0),
+ Some(1.5),
+ Some(-1.0),
+ Some(10000000000.0),
+ ]))],
+ )?;
+ let table = MemTable::try_new(schema, vec![vec![data]])?;
+
+ let mut ctx = ExecutionContext::new();
+ ctx.register_table("test", Arc::new(table))?;
+ let sql = "SELECT factorial(c1) FROM test";
Review comment:
I also recommend adding a constant here
```suggestion
let sql = "SELECT factorial(5), factorial(c1) FROM test";
```
As well as tests for what happens if you pass in a integer column
##########
File path: datafusion/src/physical_plan/math_expressions.rs
##########
@@ -102,6 +103,31 @@ math_unary_function!("ln", ln);
math_unary_function!("log2", log2);
math_unary_function!("log10", log10);
+/// factorial SQL function
+pub fn factorial(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ match &args[0] {
+ ColumnarValue::Array(array) => {
+ let x1 = array.as_any().downcast_ref::<Float64Array>();
Review comment:
I think by this point, the DataFusion coercion logic should have kicked
in and you shouldn't have to handle the cases where array is not a `Float`
array -- again I think following `sqrt` or other similar function might be
helpful
--
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]
