This is an automated email from the ASF dual-hosted git repository.
mbutrovich pushed a commit to branch df55
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/df55 by this push:
new d8cea04c2b add spark_sqrt
d8cea04c2b is described below
commit d8cea04c2b01ce12cb85cdc9459cd2b4bf35b35b
Author: Matt Butrovich <[email protected]>
AuthorDate: Wed Aug 5 09:35:01 2026 -0400
add spark_sqrt
---
native/spark-expr/src/comet_scalar_funcs.rs | 5 ++
native/spark-expr/src/lib.rs | 6 +-
native/spark-expr/src/math_funcs/mod.rs | 2 +
native/spark-expr/src/math_funcs/sqrt.rs | 90 ++++++++++++++++++++++
.../org/apache/comet/serde/QueryPlanSerde.scala | 2 +-
.../main/scala/org/apache/comet/serde/math.scala | 19 ++++-
6 files changed, 119 insertions(+), 5 deletions(-)
diff --git a/native/spark-expr/src/comet_scalar_funcs.rs
b/native/spark-expr/src/comet_scalar_funcs.rs
index 50732959c8..0173d02d35 100644
--- a/native/spark-expr/src/comet_scalar_funcs.rs
+++ b/native/spark-expr/src/comet_scalar_funcs.rs
@@ -23,6 +23,7 @@ use crate::math_funcs::checked_arithmetic::{checked_add,
checked_div, checked_mu
use crate::math_funcs::log::spark_log;
use crate::math_funcs::modulo_expr::spark_modulo;
use crate::math_funcs::pow::spark_pow;
+use crate::math_funcs::sqrt::spark_sqrt;
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,
@@ -213,6 +214,10 @@ pub fn create_comet_physical_fun_with_eval_mode(
let func = Arc::new(spark_pow);
make_comet_scalar_udf!("pow", func, without data_type)
}
+ "spark_sqrt" => {
+ let func = Arc::new(spark_sqrt);
+ make_comet_scalar_udf!("spark_sqrt", 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 42ab43361a..099f23b347 100644
--- a/native/spark-expr/src/lib.rs
+++ b/native/spark-expr/src/lib.rs
@@ -86,9 +86,9 @@ 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_pow, spark_round, spark_unhex,
spark_unscaled_value, CheckOverflow,
- DecimalRescaleCheckOverflow, NegativeExpr, NormalizeNaNAndZero,
WideDecimalBinaryExpr,
- WideDecimalOp,
+ spark_make_decimal, spark_pow, spark_round, spark_sqrt, spark_unhex,
spark_unscaled_value,
+ CheckOverflow, DecimalRescaleCheckOverflow, NegativeExpr,
NormalizeNaNAndZero,
+ WideDecimalBinaryExpr, WideDecimalOp,
};
pub use query_context::{create_query_context_map, QueryContext,
QueryContextMap};
pub use string_funcs::*;
diff --git a/native/spark-expr/src/math_funcs/mod.rs
b/native/spark-expr/src/math_funcs/mod.rs
index 56809a01fb..7a7c6a9753 100644
--- a/native/spark-expr/src/math_funcs/mod.rs
+++ b/native/spark-expr/src/math_funcs/mod.rs
@@ -26,6 +26,7 @@ pub mod modulo_expr;
mod negative;
pub(crate) mod pow;
mod round;
+pub(crate) mod sqrt;
pub(crate) mod unhex;
mod utils;
mod wide_decimal_binary_expr;
@@ -41,5 +42,6 @@ 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 sqrt::spark_sqrt;
pub use unhex::spark_unhex;
pub use wide_decimal_binary_expr::{WideDecimalBinaryExpr, WideDecimalOp};
diff --git a/native/spark-expr/src/math_funcs/sqrt.rs
b/native/spark-expr/src/math_funcs/sqrt.rs
new file mode 100644
index 0000000000..1397386eb9
--- /dev/null
+++ b/native/spark-expr/src/math_funcs/sqrt.rs
@@ -0,0 +1,90 @@
+// 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::{ArrayRef, AsArray, Float64Array};
+use arrow::datatypes::Float64Type;
+use datafusion::common::{DataFusionError, ScalarValue};
+use datafusion::physical_plan::ColumnarValue;
+use std::sync::Arc;
+
+/// Spark-compatible `sqrt`, matching `java.lang.Math.sqrt`: a negative input
produces
+/// `NaN` rather than an error, unlike DataFusion's own `sqrt`.
+pub fn spark_sqrt(args: &[ColumnarValue]) -> Result<ColumnarValue,
DataFusionError> {
+ if args.len() != 1 {
+ return Err(DataFusionError::Internal(format!(
+ "spark_sqrt requires 1 argument, got {}",
+ args.len()
+ )));
+ }
+
+ match &args[0] {
+ ColumnarValue::Array(array) => {
+ let values = array.as_primitive_opt::<Float64Type>().ok_or_else(||
{
+ DataFusionError::Internal(format!(
+ "spark_sqrt expected Float64, got {:?}",
+ array.data_type()
+ ))
+ })?;
+ let result: Float64Array = values.unary(|v| v.sqrt());
+ Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
+ }
+ ColumnarValue::Scalar(ScalarValue::Float64(v)) =>
Ok(ColumnarValue::Scalar(
+ ScalarValue::Float64(v.map(f64::sqrt)),
+ )),
+ ColumnarValue::Scalar(other) => Err(DataFusionError::Internal(format!(
+ "spark_sqrt expected Float64 scalar, got {other:?}",
+ ))),
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+ use arrow::array::{Array, Float64Array};
+
+ #[test]
+ fn test_spark_sqrt_negative_is_nan() {
+ let input = Float64Array::from(vec![Some(4.0), Some(-1.0), Some(0.0),
None]);
+ let result =
spark_sqrt(&[ColumnarValue::Array(Arc::new(input))]).unwrap();
+ let ColumnarValue::Array(result) = result else {
+ unreachable!()
+ };
+ let result = result.as_primitive::<Float64Type>();
+ assert_eq!(result.value(0), 2.0);
+ assert!(result.value(1).is_nan());
+ assert_eq!(result.value(2), 0.0);
+ assert!(result.is_null(3));
+ }
+
+ #[test]
+ fn test_spark_sqrt_scalar_negative_is_nan() {
+ let result =
+
spark_sqrt(&[ColumnarValue::Scalar(ScalarValue::Float64(Some(-1.0)))]).unwrap();
+ let ColumnarValue::Scalar(ScalarValue::Float64(Some(result))) = result
else {
+ unreachable!()
+ };
+ assert!(result.is_nan());
+ }
+
+ #[test]
+ fn test_spark_sqrt_scalar_null() {
+ let result =
spark_sqrt(&[ColumnarValue::Scalar(ScalarValue::Float64(None))]).unwrap();
+ let ColumnarValue::Scalar(ScalarValue::Float64(None)) = result else {
+ unreachable!()
+ };
+ }
+}
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 9e867aad2d..d4ac222643 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -153,7 +153,7 @@ object QueryPlanSerde extends Logging with CometExprShim
with CometTypeShim {
classOf[Signum] -> CometScalarFunction("signum"),
classOf[Sin] -> CometScalarFunction("sin"),
classOf[Sinh] -> CometScalarFunction("sinh"),
- classOf[Sqrt] -> CometScalarFunction("sqrt"),
+ classOf[Sqrt] -> CometSqrt,
classOf[Subtract] -> CometSubtract,
classOf[Tan] -> CometScalarFunction("tan"),
classOf[Tanh] -> CometScalarFunction("tanh"),
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 e70794faa3..add477c835 100644
--- a/spark/src/main/scala/org/apache/comet/serde/math.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/math.scala
@@ -19,7 +19,7 @@
package org.apache.comet.serde
-import org.apache.spark.sql.catalyst.expressions.{Abs, Add, Atan2, Attribute,
BRound, Ceil, CheckOverflow, Conv, Expression, Floor, Hex, Hypot, If,
LessThanOrEqual, Literal, Log, Log10, Log1p, Log2, Logarithm, NaNvl, Pmod, Pow,
UnaryPositive, Unhex, WidthBucket}
+import org.apache.spark.sql.catalyst.expressions.{Abs, Add, Atan2, Attribute,
BRound, Ceil, CheckOverflow, Conv, Expression, Floor, Hex, Hypot, If,
LessThanOrEqual, Literal, Log, Log10, Log1p, Log2, Logarithm, NaNvl, Pmod, Pow,
Sqrt, UnaryPositive, Unhex, WidthBucket}
import org.apache.spark.sql.types.{DecimalType, DoubleType, NumericType}
import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
@@ -217,6 +217,23 @@ object CometPow extends CometExpressionSerde[Pow] {
}
}
+// Uses a custom spark_sqrt UDF because DataFusion's own `sqrt` errors on
negative
+// input, while Spark's Sqrt (a plain wrapper around java.lang.Math.sqrt)
returns NaN.
+// spark_sqrt is a Comet-only name with no DataFusion builtin counterpart, so
the
+// return type must be set explicitly here to skip the session registry lookup
that
+// scalarFunctionExprToProto would otherwise trigger (see
CometLogarithm/spark_log).
+object CometSqrt extends CometExpressionSerde[Sqrt] {
+ override def convert(
+ expr: Sqrt,
+ inputs: Seq[Attribute],
+ binding: Boolean): Option[ExprOuterClass.Expr] = {
+ val childExpr = exprToProtoInternal(expr.child, inputs, binding)
+ val optExpr =
+ scalarFunctionExprToProtoWithReturnType("spark_sqrt", DoubleType, false,
childExpr)
+ optExprWithFallbackReason(optExpr, expr, expr.child)
+ }
+}
+
sealed trait MathExprBase {
protected def nullIfNegative(expression: Expression): Expression = {
val zero = Literal.default(expression.dataType)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
