Implementing math power function for SQL (apache#2324)

comphead · MazterQyou · commit 62b226907dd5 · 2025-05-27T17:03:19.000+04:00
* Implementing POWER function * Delete pv.yaml * Delete build-ballista-docker.sh * Delete ballista.dockerfile * aligining with latest upstream changes * Readding docker files * Formatting * Leaving only 64bit types * Adding tests, remove type conversion * fix for cast * Update functions.rs (cherry picked from commit c3c02cf) Can drop this after rebase on commit c3c02cf "Implementing math power function for SQL (apache#2324)", first released in 8.0.0 # Conflicts: # datafusion/core/src/logical_plan/mod.rs # datafusion/core/src/physical_plan/functions.rs # datafusion/core/tests/sql/functions.rs # datafusion/cube_ext/Cargo.toml # datafusion/expr/src/built_in_function.rs # datafusion/expr/src/function.rs # datafusion/proto/proto/datafusion.proto # datafusion/proto/src/from_proto.rs # datafusion/proto/src/to_proto.rs # dev/docker/ballista.dockerfile
diff --git a/datafusion/core/src/logical_plan/mod.rs b/datafusion/core/src/logical_plan/mod.rs
@@ -48,12 +48,12 @@ pub use expr::{
     count, count_distinct, create_udaf, create_udf, create_udtf, date_part, date_trunc,
     digest, exp, exprlist_to_fields, exprlist_to_fields_from_schema, floor, in_list,
     initcap, left, length, lit, lit_timestamp_nano, ln, log10, log2, lower, lpad, ltrim,
-    max, md5, min, now, now_expr, nullif, octet_length, or, pi, random, regexp_match,
-    regexp_replace, repeat, replace, reverse, right, round, rpad, rtrim, sha224, sha256,
-    sha384, sha512, signum, sin, split_part, sqrt, starts_with, strpos, substr, sum, tan,
-    to_hex, to_timestamp_micros, to_timestamp_millis, to_timestamp_seconds, translate,
-    trim, trunc, unalias, upper, when, Column, Expr, ExprSchema, GroupingSet, Like,
-    Literal,
+    max, md5, min, now, now_expr, nullif, octet_length, or, pi, power, random,
+    regexp_match, regexp_replace, repeat, replace, reverse, right, round, rpad, rtrim,
+    sha224, sha256, sha384, sha512, signum, sin, split_part, sqrt, starts_with, strpos,
+    substr, sum, tan, to_hex, to_timestamp_micros, to_timestamp_millis,
+    to_timestamp_seconds, translate, trim, trunc, unalias, upper, when, Column, Expr,
+    ExprSchema, GroupingSet, Like, Literal,
 };
 pub use expr_rewriter::{
     normalize_col, normalize_cols, replace_col, replace_col_to_expr,
diff --git a/datafusion/core/src/physical_plan/functions.rs b/datafusion/core/src/physical_plan/functions.rs
@@ -312,6 +312,10 @@ pub fn create_physical_fun(
         BuiltinScalarFunction::Sqrt => Arc::new(math_expressions::sqrt),
         BuiltinScalarFunction::Tan => Arc::new(math_expressions::tan),
         BuiltinScalarFunction::Trunc => Arc::new(math_expressions::trunc),
+        BuiltinScalarFunction::Power => {
+            Arc::new(|args| make_scalar_function(math_expressions::power)(args))
+        }
+
         BuiltinScalarFunction::Pi => Arc::new(math_expressions::pi),
         // string functions
         BuiltinScalarFunction::MakeArray => Arc::new(array_expressions::array),
diff --git a/datafusion/core/tests/sql/functions.rs b/datafusion/core/tests/sql/functions.rs
@@ -555,6 +555,132 @@ async fn case_builtin_math_expression() {
     }
 }
 
+#[tokio::test]
+async fn test_power() -> Result<()> {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("i32", DataType::Int16, true),
+        Field::new("i64", DataType::Int64, true),
+        Field::new("f32", DataType::Float32, true),
+        Field::new("f64", DataType::Float64, true),
+    ]));
+
+    let data = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int16Array::from(vec![
+                Some(2),
+                Some(5),
+                Some(0),
+                Some(-14),
+                None,
+            ])),
+            Arc::new(Int64Array::from(vec![
+                Some(2),
+                Some(5),
+                Some(0),
+                Some(-14),
+                None,
+            ])),
+            Arc::new(Float32Array::from(vec![
+                Some(1.0),
+                Some(2.5),
+                Some(0.0),
+                Some(-14.5),
+                None,
+            ])),
+            Arc::new(Float64Array::from(vec![
+                Some(1.0),
+                Some(2.5),
+                Some(0.0),
+                Some(-14.5),
+                None,
+            ])),
+        ],
+    )?;
+
+    let table = MemTable::try_new(schema, vec![vec![data]])?;
+
+    let ctx = SessionContext::new();
+    ctx.register_table("test", Arc::new(table))?;
+    let sql = r"SELECT power(i32, exp_i) as power_i32,
+                 power(i64, exp_f) as power_i64,
+                 power(f32, exp_i) as power_f32,
+                 power(f64, exp_f) as power_f64,
+                 power(2, 3) as power_int_scalar,
+                 power(2.5, 3.0) as power_float_scalar
+          FROM (select test.*, 3 as exp_i, 3.0 as exp_f from test) a";
+    let actual = execute_to_batches(&ctx, sql).await;
+    let expected = vec![
+        "+-----------+-----------+-----------+-----------+------------------+--------------------+",
+        "| power_i32 | power_i64 | power_f32 | power_f64 | power_int_scalar | power_float_scalar |",
+        "+-----------+-----------+-----------+-----------+------------------+--------------------+",
+        "| 8         | 8         | 1         | 1         | 8                | 15.625             |",
+        "| 125       | 125       | 15.625    | 15.625    | 8                | 15.625             |",
+        "| 0         | 0         | 0         | 0         | 8                | 15.625             |",
+        "| -2744     | -2744     | -3048.625 | -3048.625 | 8                | 15.625             |",
+        "|           |           |           |           | 8                | 15.625             |",
+        "+-----------+-----------+-----------+-----------+------------------+--------------------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+    //dbg!(actual[0].schema().fields());
+    assert_eq!(
+        actual[0]
+            .schema()
+            .field_with_name("power_i32")
+            .unwrap()
+            .data_type()
+            .to_owned(),
+        DataType::Int64
+    );
+    assert_eq!(
+        actual[0]
+            .schema()
+            .field_with_name("power_i64")
+            .unwrap()
+            .data_type()
+            .to_owned(),
+        DataType::Float64
+    );
+    assert_eq!(
+        actual[0]
+            .schema()
+            .field_with_name("power_f32")
+            .unwrap()
+            .data_type()
+            .to_owned(),
+        DataType::Float64
+    );
+    assert_eq!(
+        actual[0]
+            .schema()
+            .field_with_name("power_f64")
+            .unwrap()
+            .data_type()
+            .to_owned(),
+        DataType::Float64
+    );
+    assert_eq!(
+        actual[0]
+            .schema()
+            .field_with_name("power_int_scalar")
+            .unwrap()
+            .data_type()
+            .to_owned(),
+        DataType::Int64
+    );
+    assert_eq!(
+        actual[0]
+            .schema()
+            .field_with_name("power_float_scalar")
+            .unwrap()
+            .data_type()
+            .to_owned(),
+        DataType::Float64
+    );
+
+    Ok(())
+}
+
 // #[tokio::test]
 // async fn case_sensitive_identifiers_aggregates() {
 //     let ctx = SessionContext::new();
diff --git a/datafusion/expr/src/built_in_function.rs b/datafusion/expr/src/built_in_function.rs
@@ -54,6 +54,8 @@ pub enum BuiltinScalarFunction {
     Log10,
     /// log2
     Log2,
+    /// power
+    Power,
     /// pi
     Pi,
     /// round
@@ -196,6 +198,7 @@ impl BuiltinScalarFunction {
             BuiltinScalarFunction::Log => Volatility::Immutable,
             BuiltinScalarFunction::Log10 => Volatility::Immutable,
             BuiltinScalarFunction::Log2 => Volatility::Immutable,
+            BuiltinScalarFunction::Power => Volatility::Immutable,
             BuiltinScalarFunction::Pi => Volatility::Immutable,
             BuiltinScalarFunction::Round => Volatility::Immutable,
             BuiltinScalarFunction::Signum => Volatility::Immutable,
@@ -284,6 +287,7 @@ impl FromStr for BuiltinScalarFunction {
             "log" => BuiltinScalarFunction::Log,
             "log10" => BuiltinScalarFunction::Log10,
             "log2" => BuiltinScalarFunction::Log2,
+            "power" => BuiltinScalarFunction::Power,
             "pi" => BuiltinScalarFunction::Pi,
             "round" => BuiltinScalarFunction::Round,
             "signum" => BuiltinScalarFunction::Signum,
diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
@@ -266,6 +266,7 @@ unary_scalar_expr!(Log2, log2);
 unary_scalar_expr!(Log10, log10);
 unary_scalar_expr!(Ln, ln);
 unary_scalar_expr!(NullIf, nullif);
+scalar_expr!(Power, power, base, exponent);
 
 // string functions
 scalar_expr!(Ascii, ascii, string);
diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs
@@ -244,6 +244,11 @@ pub fn return_type(
             }
         }),
 
+        BuiltinScalarFunction::Power => match &input_expr_types[0] {
+            DataType::Int64 => Ok(DataType::Int64),
+            _ => Ok(DataType::Float64),
+        },
+
         BuiltinScalarFunction::Abs
         | BuiltinScalarFunction::Acos
         | BuiltinScalarFunction::Asin
@@ -550,6 +555,13 @@ pub fn signature(fun: &BuiltinScalarFunction) -> Signature {
         ),
         BuiltinScalarFunction::Pi => Signature::exact(vec![], fun.volatility()),
         BuiltinScalarFunction::Random => Signature::exact(vec![], fun.volatility()),
+        BuiltinScalarFunction::Power => Signature::one_of(
+            vec![
+                TypeSignature::Exact(vec![DataType::Int64, DataType::Int64]),
+                TypeSignature::Exact(vec![DataType::Float64, DataType::Float64]),
+            ],
+            fun.volatility(),
+        ),
         BuiltinScalarFunction::Log => Signature::one_of(
             vec![
                 TypeSignature::Exact(vec![DataType::Float64]),
diff --git a/datafusion/physical-expr/src/math_expressions.rs b/datafusion/physical-expr/src/math_expressions.rs
@@ -17,12 +17,14 @@
 
 //! Math expressions
 
-use arrow::array::{Float32Array, Float64Array};
+use arrow::array::ArrayRef;
+use arrow::array::{Float32Array, Float64Array, Int64Array};
 use arrow::datatypes::DataType;
 use datafusion_common::ScalarValue;
 use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::ColumnarValue;
 use rand::{thread_rng, Rng};
+use std::any::type_name;
 use std::iter;
 use std::sync::Arc;
 
@@ -86,6 +88,33 @@ macro_rules! math_unary_function {
     };
 }
 
+macro_rules! downcast_arg {
+    ($ARG:expr, $NAME:expr, $ARRAY_TYPE:ident) => {{
+        $ARG.as_any().downcast_ref::<$ARRAY_TYPE>().ok_or_else(|| {
+            DataFusionError::Internal(format!(
+                "could not cast {} to {}",
+                $NAME,
+                type_name::<$ARRAY_TYPE>()
+            ))
+        })?
+    }};
+}
+
+macro_rules! make_function_inputs2 {
+    ($ARG1: expr, $ARG2: expr, $NAME1:expr, $NAME2: expr, $ARRAY_TYPE:ident, $FUNC: block) => {{
+        let arg1 = downcast_arg!($ARG1, $NAME1, $ARRAY_TYPE);
+        let arg2 = downcast_arg!($ARG2, $NAME2, $ARRAY_TYPE);
+
+        arg1.iter()
+            .zip(arg2.iter())
+            .map(|(a1, a2)| match (a1, a2) {
+                (Some(a1), Some(a2)) => Some($FUNC(a1, a2.try_into().ok()?)),
+                _ => None,
+            })
+            .collect::<$ARRAY_TYPE>()
+    }};
+}
+
 math_unary_function!("sqrt", sqrt);
 math_unary_function!("sin", sin);
 math_unary_function!("cos", cos);
@@ -131,6 +160,33 @@ pub fn random(args: &[ColumnarValue]) -> Result<ColumnarValue> {
     Ok(ColumnarValue::Array(Arc::new(array)))
 }
 
+pub fn power(args: &[ArrayRef]) -> Result<ArrayRef> {
+    match args[0].data_type() {
+        DataType::Float64 => Ok(Arc::new(make_function_inputs2!(
+            &args[0],
+            &args[1],
+            "base",
+            "exponent",
+            Float64Array,
+            { f64::powf }
+        )) as ArrayRef),
+
+        DataType::Int64 => Ok(Arc::new(make_function_inputs2!(
+            &args[0],
+            &args[1],
+            "base",
+            "exponent",
+            Int64Array,
+            { i64::pow }
+        )) as ArrayRef),
+
+        other => Err(DataFusionError::Internal(format!(
+            "Unsupported data type {:?} for function power",
+            other
+        ))),
+    }
+}
+
 #[cfg(test)]
 mod tests {
 
diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
@@ -190,6 +190,7 @@ enum ScalarFunction {
   Upper=62;
   Coalesce=63;
   // Upstream
+  Power=64;
   CurrentDate=70;
   Pi=80;
   // Cubesql
diff --git a/datafusion/proto/src/from_proto.rs b/datafusion/proto/src/from_proto.rs
@@ -25,9 +25,9 @@ use datafusion::{
     logical_plan::{
         abs, acos, ascii, asin, atan, ceil, character_length, chr, concat_expr,
         concat_ws_expr, cos, digest, exp, floor, left, ln, log10, log2, now_expr, nullif,
-        pi, random, regexp_replace, repeat, replace, reverse, right, round, signum, sin,
-        split_part, sqrt, starts_with, strpos, substr, tan, to_hex, to_timestamp_micros,
-        to_timestamp_millis, to_timestamp_seconds, translate, trunc,
+        pi, power, random, regexp_replace, repeat, replace, reverse, right, round,
+        signum, sin, split_part, sqrt, starts_with, strpos, substr, tan, to_hex,
+        to_timestamp_micros, to_timestamp_millis, to_timestamp_seconds, translate, trunc,
         window_frames::{WindowFrame, WindowFrameBound, WindowFrameUnits},
         Column, DFField, DFSchema, DFSchemaRef, Expr, Like, Operator,
     },
@@ -430,6 +430,7 @@ impl From<&protobuf::ScalarFunction> for BuiltinScalarFunction {
             ScalarFunction::Translate => Self::Translate,
             ScalarFunction::RegexpMatch => Self::RegexpMatch,
             ScalarFunction::Coalesce => Self::Coalesce,
+            ScalarFunction::Power => Self::Power,
             ScalarFunction::Pi => Self::Pi,
             // Cube SQL
             ScalarFunction::UtcTimestamp => Self::UtcTimestamp,
@@ -1232,6 +1233,10 @@ pub fn parse_expr(
                         .map(|expr| parse_expr(expr, registry))
                         .collect::<Result<Vec<_>, _>>()?,
                 )),
+                ScalarFunction::Power => Ok(power(
+                    parse_expr(&args[0], registry)?,
+                    parse_expr(&args[1], registry)?,
+                )),
                 ScalarFunction::Pi => Ok(pi()),
                 _ => Err(proto_error(
                     "Protobuf deserialization error: Unsupported scalar function",
diff --git a/datafusion/proto/src/to_proto.rs b/datafusion/proto/src/to_proto.rs
@@ -1079,6 +1079,7 @@ impl TryFrom<&BuiltinScalarFunction> for protobuf::ScalarFunction {
             BuiltinScalarFunction::Translate => Self::Translate,
             BuiltinScalarFunction::RegexpMatch => Self::RegexpMatch,
             BuiltinScalarFunction::Coalesce => Self::Coalesce,
+            BuiltinScalarFunction::Power => Self::Power,
             BuiltinScalarFunction::Pi => Self::Pi,
             // Cube SQL
             BuiltinScalarFunction::UtcTimestamp => Self::UtcTimestamp,
