diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 8e913dabab..925f6cb335 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -22,6 +22,7 @@ use crate::math_funcs::abs::abs; use crate::math_funcs::checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub}; use crate::math_funcs::log::spark_log; use crate::math_funcs::modulo_expr::spark_modulo; +use crate::math_funcs::pow::spark_pow; use crate::{ spark_ceil, spark_day_name, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding, @@ -193,6 +194,10 @@ pub fn create_comet_physical_fun_with_eval_mode( let func = Arc::new(spark_log); make_comet_scalar_udf!("spark_log", func, without data_type) } + "pow" => { + let func = Arc::new(spark_pow); + make_comet_scalar_udf!("pow", func, without data_type) + } "base64" => { let func = Arc::new(crate::string_funcs::spark_base64); make_comet_scalar_udf!("base64", func, without data_type) diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index 40422cde07..c30eaff9c9 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -86,7 +86,7 @@ pub use json_funcs::{FromJson, ToJson}; pub use math_funcs::{ checked_add, checked_div, checked_mul, checked_sub, create_modulo_expr, create_negate_expr, spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_log, - spark_make_decimal, spark_round, spark_unhex, spark_unscaled_value, CheckOverflow, + spark_make_decimal, spark_pow, spark_round, spark_unhex, spark_unscaled_value, CheckOverflow, DecimalRescaleCheckOverflow, NegativeExpr, NormalizeNaNAndZero, WideDecimalBinaryExpr, WideDecimalOp, }; diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs index 51bd016a2a..56809a01fb 100644 --- a/native/spark-expr/src/math_funcs/mod.rs +++ b/native/spark-expr/src/math_funcs/mod.rs @@ -24,6 +24,7 @@ pub mod internal; pub(crate) mod log; pub mod modulo_expr; mod negative; +pub(crate) mod pow; mod round; pub(crate) mod unhex; mod utils; @@ -38,6 +39,7 @@ pub use internal::*; pub use log::spark_log; pub use modulo_expr::create_modulo_expr; pub use negative::{create_negate_expr, NegativeExpr}; +pub use pow::spark_pow; pub use round::spark_round; pub use unhex::spark_unhex; pub use wide_decimal_binary_expr::{WideDecimalBinaryExpr, WideDecimalOp}; diff --git a/native/spark-expr/src/math_funcs/pow.rs b/native/spark-expr/src/math_funcs/pow.rs new file mode 100644 index 0000000000..61168f7d18 --- /dev/null +++ b/native/spark-expr/src/math_funcs/pow.rs @@ -0,0 +1,314 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::Float64Array; +use datafusion::common::{DataFusionError, ScalarValue}; +use datafusion::physical_plan::ColumnarValue; +use std::sync::Arc; + +/// Spark-compatible scalar power matching Java's `Math.pow`. +/// +/// Rust's `f64::powf` follows C99 `pow` semantics, which agree with `Math.pow` on almost every +/// input (including `pow(0, -1) == Infinity`, the signed-zero rules, and the infinite-base rules) +/// but diverge in one place: when `|base| == 1` and the exponent is infinite or `NaN`, C99 `pow` +/// returns `1` whereas `Math.pow` returns `NaN`. Special-case that so the native result matches +/// Spark exactly. `Math.pow` still returns `1` for `pow(x, ±0)` even when the base is non-finite, +/// but that is a finite exponent and is left to `powf`. +#[inline] +fn spark_powf(base: f64, exp: f64) -> f64 { + if base.abs() == 1.0 && !exp.is_finite() { + return f64::NAN; + } + base.powf(exp) +} + +/// Spark-compatible power: `pow(base, exponent)`. +/// +/// Matches Spark's `Pow` expression, which delegates to Java's `Math.pow`, via [`spark_powf`]. +/// Unlike DataFusion's `power`, `pow(0, -1)` returns `Infinity` rather than erroring. Only null +/// inputs produce null; otherwise every result is the `spark_powf` value. +pub fn spark_pow(args: &[ColumnarValue]) -> Result { + if args.len() != 2 { + return Err(DataFusionError::Internal(format!( + "spark_pow requires 2 arguments, got {}", + args.len() + ))); + } + + fn as_f64_array( + value: &Arc, + ) -> Result<&Float64Array, DataFusionError> { + value + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal(format!( + "spark_pow expected Float64, got {:?}", + value.data_type() + )) + }) + } + + fn as_f64_scalar(scalar: &ScalarValue) -> Result, DataFusionError> { + match scalar { + ScalarValue::Float64(v) => Ok(*v), + _ => Err(DataFusionError::Internal(format!( + "spark_pow expected Float64 scalar, got {scalar:?}", + ))), + } + } + + match (&args[0], &args[1]) { + (ColumnarValue::Array(base_arr), ColumnarValue::Array(exp_arr)) => { + let bases = as_f64_array(base_arr)?; + let exps = as_f64_array(exp_arr)?; + let result: Float64Array = bases + .iter() + .zip(exps.iter()) + .map(|(b, e)| match (b, e) { + (Some(base), Some(exp)) => Some(spark_powf(base, exp)), + _ => None, + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(result))) + } + (ColumnarValue::Scalar(base_scalar), ColumnarValue::Array(exp_arr)) => { + let exps = as_f64_array(exp_arr)?; + let result: Float64Array = match as_f64_scalar(base_scalar)? { + Some(base) => exps + .iter() + .map(|e| e.map(|exp| spark_powf(base, exp))) + .collect(), + None => Float64Array::new_null(exp_arr.len()), + }; + Ok(ColumnarValue::Array(Arc::new(result))) + } + (ColumnarValue::Array(base_arr), ColumnarValue::Scalar(exp_scalar)) => { + let bases = as_f64_array(base_arr)?; + let result: Float64Array = match as_f64_scalar(exp_scalar)? { + Some(exp) => bases + .iter() + .map(|b| b.map(|base| spark_powf(base, exp))) + .collect(), + None => Float64Array::new_null(base_arr.len()), + }; + Ok(ColumnarValue::Array(Arc::new(result))) + } + (ColumnarValue::Scalar(base_scalar), ColumnarValue::Scalar(exp_scalar)) => { + let result = match (as_f64_scalar(base_scalar)?, as_f64_scalar(exp_scalar)?) { + (Some(base), Some(exp)) => ScalarValue::Float64(Some(spark_powf(base, exp))), + _ => ScalarValue::Float64(None), + }; + Ok(ColumnarValue::Scalar(result)) + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use arrow::array::Array; + + #[test] + fn test_spark_pow_basic() { + let bases = Float64Array::from(vec![2.0, 2.0, -1.0]); + let exps = Float64Array::from(vec![3.0, -1.0, 2.0]); + let result = spark_pow(&[ + ColumnarValue::Array(Arc::new(bases)), + ColumnarValue::Array(Arc::new(exps)), + ]) + .unwrap(); + if let ColumnarValue::Array(arr) = result { + let arr = arr.as_any().downcast_ref::().unwrap(); + assert!((arr.value(0) - 8.0).abs() < 1e-10); + assert!((arr.value(1) - 0.5).abs() < 1e-10); + assert!((arr.value(2) - 1.0).abs() < 1e-10); + } else { + panic!("expected array result"); + } + } + + /// Evaluate `spark_pow` over paired base/exponent columns and return the result array. + fn eval_pairs(bases: Vec, exps: Vec) -> Float64Array { + let result = spark_pow(&[ + ColumnarValue::Array(Arc::new(Float64Array::from(bases))), + ColumnarValue::Array(Arc::new(Float64Array::from(exps))), + ]) + .unwrap(); + match result { + ColumnarValue::Array(arr) => { + arr.as_any().downcast_ref::().unwrap().clone() + } + _ => panic!("expected array result"), + } + } + + /// Assert a value is a negative zero (distinct from +0.0, which compares equal under `==`). + fn assert_negative_zero(v: f64) { + assert_eq!(v, 0.0, "expected zero, got {v}"); + assert!(v.is_sign_negative(), "expected negative zero, got +0.0"); + } + + #[test] + fn test_spark_pow_zero_negative_exp_is_infinity() { + // Spark/Java Math.pow(0, -1) == +Infinity (DataFusion's power errors here instead). + let arr = eval_pairs(vec![0.0], vec![-1.0]); + assert_eq!(arr.value(0), f64::INFINITY); + } + + #[test] + fn test_spark_pow_abs_one_nonfinite_exp_is_nan() { + // Java Math.pow returns NaN when |base| == 1 and the exponent is infinite or NaN, whereas + // C99 pow / Rust powf return 1. spark_pow must match Spark and return NaN. + let arr = eval_pairs( + vec![1.0, -1.0, 1.0, -1.0, 1.0, -1.0], + vec![ + f64::INFINITY, + f64::INFINITY, + f64::NEG_INFINITY, + f64::NEG_INFINITY, + f64::NAN, + f64::NAN, + ], + ); + for i in 0..arr.len() { + assert!( + arr.value(i).is_nan(), + "row {i} expected NaN, got {}", + arr.value(i) + ); + } + // A finite exponent on |base| == 1 is unaffected: pow(1, 0) == 1, pow(-1, 3) == -1. + let finite = eval_pairs(vec![1.0, -1.0], vec![0.0, 3.0]); + assert_eq!(finite.value(0), 1.0); + assert_eq!(finite.value(1), -1.0); + } + + #[test] + fn test_spark_pow_negative_zero_base() { + // Signed-zero rules from Math.pow: odd/even and sign of exponent select sign/infinity. + let arr = eval_pairs(vec![-0.0, -0.0, -0.0, -0.0], vec![-1.0, -2.0, 3.0, 2.0]); + assert_eq!(arr.value(0), f64::NEG_INFINITY); // (-0)^-1 + assert_eq!(arr.value(1), f64::INFINITY); // (-0)^-2 + assert_negative_zero(arr.value(2)); // (-0)^3 == -0.0 + assert_eq!(arr.value(3), 0.0); // (-0)^2 == +0.0 + assert!(arr.value(3).is_sign_positive()); + } + + #[test] + fn test_spark_pow_infinite_base_and_exp() { + let arr = eval_pairs( + vec![ + f64::NEG_INFINITY, + f64::NEG_INFINITY, + f64::NEG_INFINITY, + 2.0, + 0.5, + ], + vec![2.0, 3.0, -1.0, f64::NEG_INFINITY, f64::NEG_INFINITY], + ); + assert_eq!(arr.value(0), f64::INFINITY); // (-inf)^2 + assert_eq!(arr.value(1), f64::NEG_INFINITY); // (-inf)^3 + assert_negative_zero(arr.value(2)); // (-inf)^-1 == -0.0 + assert_eq!(arr.value(3), 0.0); // 2^-inf == +0.0 + assert_eq!(arr.value(4), f64::INFINITY); // 0.5^-inf == +inf + } + + #[test] + fn test_spark_pow_subnormal() { + // Smallest positive subnormal (Double.MIN_VALUE). Squaring underflows to +0.0; raising a + // finite base to a subnormal exponent rounds to 1.0. Both match Spark. + let min_subnormal = f64::from_bits(1); + let arr = eval_pairs(vec![min_subnormal, 2.0], vec![2.0, min_subnormal]); + assert_eq!(arr.value(0), 0.0); + assert_eq!(arr.value(1), 1.0); + } + + #[test] + fn test_spark_pow_scalar_abs_one_nonfinite_exp_is_nan() { + // The |base| == 1 fix must apply on the scalar-base and scalar-exponent paths too. + let scalar_base = spark_pow(&[ + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0))), + ColumnarValue::Array(Arc::new(Float64Array::from(vec![f64::INFINITY]))), + ]) + .unwrap(); + if let ColumnarValue::Array(arr) = scalar_base { + let arr = arr.as_any().downcast_ref::().unwrap(); + assert!(arr.value(0).is_nan()); + } else { + panic!("expected array result"); + } + + let scalar_exp = spark_pow(&[ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![-1.0]))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::INFINITY))), + ]) + .unwrap(); + if let ColumnarValue::Array(arr) = scalar_exp { + let arr = arr.as_any().downcast_ref::().unwrap(); + assert!(arr.value(0).is_nan()); + } else { + panic!("expected array result"); + } + + let both_scalar = spark_pow(&[ + ColumnarValue::Scalar(ScalarValue::Float64(Some(-1.0))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::NAN))), + ]) + .unwrap(); + if let ColumnarValue::Scalar(ScalarValue::Float64(Some(v))) = both_scalar { + assert!(v.is_nan()); + } else { + panic!("expected scalar float64 result"); + } + } + + #[test] + fn test_spark_pow_null_propagation() { + let bases = Float64Array::from(vec![Some(2.0), None]); + let exps = Float64Array::from(vec![None, Some(2.0)]); + let result = spark_pow(&[ + ColumnarValue::Array(Arc::new(bases)), + ColumnarValue::Array(Arc::new(exps)), + ]) + .unwrap(); + if let ColumnarValue::Array(arr) = result { + let arr = arr.as_any().downcast_ref::().unwrap(); + assert!(arr.is_null(0)); + assert!(arr.is_null(1)); + } else { + panic!("expected array result"); + } + } + + #[test] + fn test_spark_pow_scalar_base() { + let exps = Float64Array::from(vec![Some(3.0), None]); + let result = spark_pow(&[ + ColumnarValue::Scalar(ScalarValue::Float64(Some(2.0))), + ColumnarValue::Array(Arc::new(exps)), + ]) + .unwrap(); + if let ColumnarValue::Array(arr) = result { + let arr = arr.as_any().downcast_ref::().unwrap(); + assert!((arr.value(0) - 8.0).abs() < 1e-10); + assert!(arr.is_null(1)); + } else { + panic!("expected array result"); + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/serde/math.scala b/spark/src/main/scala/org/apache/comet/serde/math.scala index 0ec2b5cd9e..e70794faa3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/math.scala +++ b/spark/src/main/scala/org/apache/comet/serde/math.scala @@ -206,14 +206,6 @@ object CometAbs extends CometExpressionSerde[Abs] with MathExprBase { object CometPow extends CometExpressionSerde[Pow] { - // https://github.com/apache/datafusion/issues/22598 - val unsupportedReason: String = "Power has correctness issues" - - override def getUnsupportedReasons(): Seq[String] = Seq(unsupportedReason) - - override def getSupportLevel(expr: Pow): SupportLevel = - Unsupported(Some(unsupportedReason)) - override def convert( expr: Pow, inputs: Seq[Attribute], diff --git a/spark/src/test/resources/sql-tests/expressions/math/pow.sql b/spark/src/test/resources/sql-tests/expressions/math/pow.sql index f1111167ba..f6baddffa9 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/pow.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/pow.sql @@ -15,27 +15,42 @@ -- specific language governing permissions and limitations -- under the License. +-- pow runs natively and matches Spark exactly, including edge cases such as pow(0, -1) = Infinity. + statement CREATE TABLE test_pow(base double, exp double) USING parquet +-- Rows cover: ordinary values, nulls, signed zero, positive/negative infinity in both base and +-- exponent, subnormal values, and the |base| == 1 with non-finite exponent case where Java's +-- Math.pow returns NaN but C99 pow (Rust powf) returns 1. statement -INSERT INTO test_pow VALUES (0.0, -1), (2.0, 3.0), (0.0, 0.0), (-1.0, 2.0), (-1.0, 0.5), (2.0, -1.0), (NULL, 2.0), (2.0, NULL), (cast('NaN' as double), 2.0), (cast('Infinity' as double), 2.0), (2.0, cast('Infinity' as double)) +INSERT INTO test_pow VALUES + (0.0, -1.0), (2.0, 3.0), (0.0, 0.0), (-1.0, 2.0), (-1.0, 0.5), (2.0, -1.0), + (NULL, 2.0), (2.0, NULL), + (cast('NaN' as double), 2.0), (cast('Infinity' as double), 2.0), (2.0, cast('Infinity' as double)), + (cast('-0.0' as double), -1.0), (cast('-0.0' as double), -2.0), (cast('-0.0' as double), 3.0), (cast('-0.0' as double), 2.0), + (cast('-Infinity' as double), 2.0), (cast('-Infinity' as double), 3.0), (cast('-Infinity' as double), -1.0), + (2.0, cast('-Infinity' as double)), (0.5, cast('-Infinity' as double)), + (1.0, cast('Infinity' as double)), (-1.0, cast('Infinity' as double)), (1.0, cast('-Infinity' as double)), + (1.0, cast('NaN' as double)), (-1.0, cast('NaN' as double)), + (cast('4.9E-324' as double), 2.0), (2.0, cast('4.9E-324' as double)) --- query tolerance=1e-6 -query expect_fallback(Power has correctness issues) +-- Every pair above yields an exact double, so use exact comparison (no tolerance) to assert the +-- NaN and signed-Infinity edge cases rather than silently skipping them: a tolerance comparison +-- makes no assertion at all for NaN and ignores the sign of Infinity. Exact comparison is +-- payload-insensitive for NaN (the test framework canonicalizes it) but not for signed zero, so +-- the -0.0 results are asserted exactly. +query SELECT pow(base, exp) FROM test_pow -- column + literal --- query tolerance=1e-6 -query expect_fallback(Power has correctness issues) +query tolerance=1e-6 SELECT pow(base, 2.0) FROM test_pow -- literal + column --- query tolerance=1e-6 -query expect_fallback(Power has correctness issues) +query tolerance=1e-6 SELECT pow(2.0, exp) FROM test_pow -- literal + literal --- query tolerance=1e-6 -query expect_fallback(Power has correctness issues) +query tolerance=1e-6 SELECT pow(2.0, 3.0), pow(0.0, 0.0), pow(-1.0, 2.0), pow(NULL, 2.0) diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 9a59c9c2eb..a1082850c1 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -429,8 +429,17 @@ abstract class CometTestBase case s: java.lang.Short => s.shortValue case i: java.lang.Integer => i.intValue case l: java.lang.Long => l.longValue - case f: java.lang.Float => f.floatValue - case d: java.lang.Double => d.doubleValue + // `QueryTest.compare` compares floating point values by raw bits so that 0.0 and -0.0 are + // distinguished. A side effect is that it also distinguishes NaN payloads, which are not + // specified: `Math.pow` and friends may return any NaN, and the payload the JVM produces + // differs across architectures, so a native kernel and Spark can return NaNs that print + // identically but compare unequal. Spark fixed this in `compare` itself for 4.1.2 and 4.2.0 + // ("in some hardware NaN can be represented with different bits, so first check for it"), but + // 3.4, 3.5, 4.0 and 4.1.1 still compare raw bits. Canonicalize NaN here so every supported + // version asserts NaN-ness without asserting the payload. Signed zero is deliberately left + // alone: distinguishing it is the reason the raw-bit comparison exists. + case f: java.lang.Float => if (f.isNaN) Float.NaN else f.floatValue + case d: java.lang.Double => if (d.isNaN) Double.NaN else d.doubleValue case x => x }