datafusion_functions/math/
power.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
18//! Math function: `power()`.
19use std::any::Any;
20use std::sync::Arc;
21
22use super::log::LogFunc;
23
24use arrow::array::{ArrayRef, AsArray, Int64Array};
25use arrow::datatypes::{ArrowNativeTypeOp, DataType, Float64Type};
26use datafusion_common::{
27    arrow_datafusion_err, exec_datafusion_err, exec_err, internal_datafusion_err,
28    plan_datafusion_err, DataFusionError, Result, ScalarValue,
29};
30use datafusion_expr::expr::ScalarFunction;
31use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
32use datafusion_expr::{
33    ColumnarValue, Documentation, Expr, ScalarFunctionArgs, ScalarUDF, TypeSignature,
34};
35use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
36use datafusion_macros::user_doc;
37
38#[user_doc(
39    doc_section(label = "Math Functions"),
40    description = "Returns a base expression raised to the power of an exponent.",
41    syntax_example = "power(base, exponent)",
42    standard_argument(name = "base", prefix = "Numeric"),
43    standard_argument(name = "exponent", prefix = "Exponent numeric")
44)]
45#[derive(Debug)]
46pub struct PowerFunc {
47    signature: Signature,
48    aliases: Vec<String>,
49}
50
51impl Default for PowerFunc {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57impl PowerFunc {
58    pub fn new() -> Self {
59        use DataType::*;
60        Self {
61            signature: Signature::one_of(
62                vec![
63                    TypeSignature::Exact(vec![Int64, Int64]),
64                    TypeSignature::Exact(vec![Float64, Float64]),
65                ],
66                Volatility::Immutable,
67            ),
68            aliases: vec![String::from("pow")],
69        }
70    }
71}
72
73impl ScalarUDFImpl for PowerFunc {
74    fn as_any(&self) -> &dyn Any {
75        self
76    }
77    fn name(&self) -> &str {
78        "power"
79    }
80
81    fn signature(&self) -> &Signature {
82        &self.signature
83    }
84
85    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
86        match arg_types[0] {
87            DataType::Int64 => Ok(DataType::Int64),
88            _ => Ok(DataType::Float64),
89        }
90    }
91
92    fn aliases(&self) -> &[String] {
93        &self.aliases
94    }
95
96    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
97        let args = ColumnarValue::values_to_arrays(&args.args)?;
98
99        let arr: ArrayRef = match args[0].data_type() {
100            DataType::Float64 => {
101                let bases = args[0].as_primitive::<Float64Type>();
102                let exponents = args[1].as_primitive::<Float64Type>();
103                let result = arrow::compute::binary::<_, _, _, Float64Type>(
104                    bases,
105                    exponents,
106                    f64::powf,
107                )?;
108                Arc::new(result) as _
109            }
110            DataType::Int64 => {
111                let bases = downcast_named_arg!(&args[0], "base", Int64Array);
112                let exponents = downcast_named_arg!(&args[1], "exponent", Int64Array);
113                bases
114                    .iter()
115                    .zip(exponents.iter())
116                    .map(|(base, exp)| match (base, exp) {
117                        (Some(base), Some(exp)) => Ok(Some(base.pow_checked(
118                            exp.try_into().map_err(|_| {
119                                exec_datafusion_err!(
120                                    "Can't use negative exponents: {exp} in integer computation, please use Float."
121                                )
122                            })?,
123                        ).map_err(|e| arrow_datafusion_err!(e))?)),
124                        _ => Ok(None),
125                    })
126                    .collect::<Result<Int64Array>>()
127                    .map(Arc::new)? as _
128            }
129
130            other => {
131                return exec_err!(
132                    "Unsupported data type {other:?} for function {}",
133                    self.name()
134                )
135            }
136        };
137
138        Ok(ColumnarValue::Array(arr))
139    }
140
141    /// Simplify the `power` function by the relevant rules:
142    /// 1. Power(a, 0) ===> 0
143    /// 2. Power(a, 1) ===> a
144    /// 3. Power(a, Log(a, b)) ===> b
145    fn simplify(
146        &self,
147        mut args: Vec<Expr>,
148        info: &dyn SimplifyInfo,
149    ) -> Result<ExprSimplifyResult> {
150        let exponent = args.pop().ok_or_else(|| {
151            plan_datafusion_err!("Expected power to have 2 arguments, got 0")
152        })?;
153        let base = args.pop().ok_or_else(|| {
154            plan_datafusion_err!("Expected power to have 2 arguments, got 1")
155        })?;
156
157        let exponent_type = info.get_data_type(&exponent)?;
158        match exponent {
159            Expr::Literal(value) if value == ScalarValue::new_zero(&exponent_type)? => {
160                Ok(ExprSimplifyResult::Simplified(Expr::Literal(
161                    ScalarValue::new_one(&info.get_data_type(&base)?)?,
162                )))
163            }
164            Expr::Literal(value) if value == ScalarValue::new_one(&exponent_type)? => {
165                Ok(ExprSimplifyResult::Simplified(base))
166            }
167            Expr::ScalarFunction(ScalarFunction { func, mut args })
168                if is_log(&func) && args.len() == 2 && base == args[0] =>
169            {
170                let b = args.pop().unwrap(); // length checked above
171                Ok(ExprSimplifyResult::Simplified(b))
172            }
173            _ => Ok(ExprSimplifyResult::Original(vec![base, exponent])),
174        }
175    }
176
177    fn documentation(&self) -> Option<&Documentation> {
178        self.doc()
179    }
180}
181
182/// Return true if this function call is a call to `Log`
183fn is_log(func: &ScalarUDF) -> bool {
184    func.inner().as_any().downcast_ref::<LogFunc>().is_some()
185}
186
187#[cfg(test)]
188mod tests {
189    use arrow::array::Float64Array;
190    use datafusion_common::cast::{as_float64_array, as_int64_array};
191
192    use super::*;
193
194    #[test]
195    fn test_power_f64() {
196        let args = ScalarFunctionArgs {
197            args: vec![
198                ColumnarValue::Array(Arc::new(Float64Array::from(vec![
199                    2.0, 2.0, 3.0, 5.0,
200                ]))), // base
201                ColumnarValue::Array(Arc::new(Float64Array::from(vec![
202                    3.0, 2.0, 4.0, 4.0,
203                ]))), // exponent
204            ],
205            number_rows: 4,
206            return_type: &DataType::Float64,
207        };
208        let result = PowerFunc::new()
209            .invoke_with_args(args)
210            .expect("failed to initialize function power");
211
212        match result {
213            ColumnarValue::Array(arr) => {
214                let floats = as_float64_array(&arr)
215                    .expect("failed to convert result to a Float64Array");
216                assert_eq!(floats.len(), 4);
217                assert_eq!(floats.value(0), 8.0);
218                assert_eq!(floats.value(1), 4.0);
219                assert_eq!(floats.value(2), 81.0);
220                assert_eq!(floats.value(3), 625.0);
221            }
222            ColumnarValue::Scalar(_) => {
223                panic!("Expected an array value")
224            }
225        }
226    }
227
228    #[test]
229    fn test_power_i64() {
230        let args = ScalarFunctionArgs {
231            args: vec![
232                ColumnarValue::Array(Arc::new(Int64Array::from(vec![2, 2, 3, 5]))), // base
233                ColumnarValue::Array(Arc::new(Int64Array::from(vec![3, 2, 4, 4]))), // exponent
234            ],
235            number_rows: 4,
236            return_type: &DataType::Int64,
237        };
238        let result = PowerFunc::new()
239            .invoke_with_args(args)
240            .expect("failed to initialize function power");
241
242        match result {
243            ColumnarValue::Array(arr) => {
244                let ints = as_int64_array(&arr)
245                    .expect("failed to convert result to a Int64Array");
246
247                assert_eq!(ints.len(), 4);
248                assert_eq!(ints.value(0), 8);
249                assert_eq!(ints.value(1), 4);
250                assert_eq!(ints.value(2), 81);
251                assert_eq!(ints.value(3), 625);
252            }
253            ColumnarValue::Scalar(_) => {
254                panic!("Expected an array value")
255            }
256        }
257    }
258}
datafusion_functions/math/power.rs

datafusion_functions/math/
power.rs
