datafusion_functions/math/
factorial.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::{
19    array::{ArrayRef, Int64Array},
20    error::ArrowError,
21};
22use std::any::Any;
23use std::sync::Arc;
24
25use arrow::datatypes::DataType;
26use arrow::datatypes::DataType::Int64;
27
28use crate::utils::make_scalar_function;
29use datafusion_common::{arrow_datafusion_err, exec_err, DataFusionError, Result};
30use datafusion_expr::{
31    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    Volatility,
33};
34use datafusion_macros::user_doc;
35
36#[user_doc(
37    doc_section(label = "Math Functions"),
38    description = "Factorial. Returns 1 if value is less than 2.",
39    syntax_example = "factorial(numeric_expression)",
40    sql_example = r#"```sql
41> SELECT factorial(5);
42+---------------+
43| factorial(5)  |
44+---------------+
45| 120           |
46+---------------+
47```"#,
48    standard_argument(name = "numeric_expression", prefix = "Numeric")
49)]
50#[derive(Debug, PartialEq, Eq, Hash)]
51pub struct FactorialFunc {
52    signature: Signature,
53}
54
55impl Default for FactorialFunc {
56    fn default() -> Self {
57        FactorialFunc::new()
58    }
59}
60
61impl FactorialFunc {
62    pub fn new() -> Self {
63        Self {
64            signature: Signature::uniform(1, vec![Int64], Volatility::Immutable),
65        }
66    }
67}
68
69impl ScalarUDFImpl for FactorialFunc {
70    fn as_any(&self) -> &dyn Any {
71        self
72    }
73
74    fn name(&self) -> &str {
75        "factorial"
76    }
77
78    fn signature(&self) -> &Signature {
79        &self.signature
80    }
81
82    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
83        Ok(Int64)
84    }
85
86    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
87        make_scalar_function(factorial, vec![])(&args.args)
88    }
89
90    fn documentation(&self) -> Option<&Documentation> {
91        self.doc()
92    }
93}
94
95/// Factorial SQL function
96fn factorial(args: &[ArrayRef]) -> Result<ArrayRef> {
97    match args[0].data_type() {
98        Int64 => {
99            let arg = downcast_named_arg!((&args[0]), "value", Int64Array);
100            Ok(arg
101                .iter()
102                .map(|a| match a {
103                    Some(a) => (2..=a)
104                        .try_fold(1i64, i64::checked_mul)
105                        .ok_or_else(|| {
106                            arrow_datafusion_err!(ArrowError::ComputeError(format!(
107                                "Overflow happened on FACTORIAL({a})"
108                            )))
109                        })
110                        .map(Some),
111                    _ => Ok(None),
112                })
113                .collect::<Result<Int64Array>>()
114                .map(Arc::new)? as ArrayRef)
115        }
116        other => exec_err!("Unsupported data type {other:?} for function factorial."),
117    }
118}
119
120#[cfg(test)]
121mod test {
122
123    use datafusion_common::cast::as_int64_array;
124
125    use super::*;
126
127    #[test]
128    fn test_factorial_i64() {
129        let args: Vec<ArrayRef> = vec![
130            Arc::new(Int64Array::from(vec![0, 1, 2, 4])), // input
131        ];
132
133        let result = factorial(&args).expect("failed to initialize function factorial");
134        let ints =
135            as_int64_array(&result).expect("failed to initialize function factorial");
136
137        let expected = Int64Array::from(vec![1, 1, 2, 24]);
138
139        assert_eq!(ints, &expected);
140    }
141}
datafusion_functions/math/factorial.rs

datafusion_functions/math/
factorial.rs
