From 8cae008cb0d0875be6bbccf432c06c3cf04f4f74 Mon Sep 17 00:00:00 2001 From: KarpagamKarthikeyan Date: Tue, 21 Jul 2026 12:42:25 -0700 Subject: [PATCH 1/2] feat: add Spark-compatible hypot function --- datafusion/spark/src/function/math/hypot.rs | 90 +++++++++++++++++++ datafusion/spark/src/function/math/mod.rs | 4 + .../test_files/spark/math/hypot.slt | 46 +++++++++- 3 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 datafusion/spark/src/function/math/hypot.rs diff --git a/datafusion/spark/src/function/math/hypot.rs b/datafusion/spark/src/function/math/hypot.rs new file mode 100644 index 0000000000000..2f88a3fd9559d --- /dev/null +++ b/datafusion/spark/src/function/math/hypot.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 std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, Float64Array}; +use arrow::compute::kernels::arity::binary; +use arrow::datatypes::{DataType, Float64Type}; +use datafusion_common::Result; +use datafusion_common::utils::take_function_args; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; + +/// Spark-compatible `hypot` function. +/// +/// +/// +/// Returns `sqrt(expr1^2 + expr2^2)` computed without intermediate overflow or +/// underflow, matching Spark's use of `java.lang.Math.hypot`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkHypot { + signature: Signature, +} + +impl Default for SparkHypot { + fn default() -> Self { + Self::new() + } +} + +impl SparkHypot { + pub fn new() -> Self { + Self { + // Spark only defines hypot over doubles; `exact` makes coercion + // guarantee both inputs are Float64 before `invoke` runs. + signature: Signature::exact( + vec![DataType::Float64, DataType::Float64], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkHypot { + fn name(&self) -> &str { + "hypot" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let num_rows = args.number_rows; + let [x, y] = take_function_args(self.name(), args.args)?; + + // Broadcast scalars to arrays so one path covers every combination. + let x = x.to_array(num_rows)?; + let y = y.to_array(num_rows)?; + + // Safe: the `exact` signature guarantees Float64 inputs. + let x = x.as_primitive::(); + let y = y.as_primitive::(); + + // `binary` applies the op element-wise and returns NULL when either + // input is NULL — matching Spark's null semantics. + let result: Float64Array = binary(x, y, |a, b| a.hypot(b))?; + + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} diff --git a/datafusion/spark/src/function/math/mod.rs b/datafusion/spark/src/function/math/mod.rs index 0079ef0fc97cd..fb57b536f26ec 100644 --- a/datafusion/spark/src/function/math/mod.rs +++ b/datafusion/spark/src/function/math/mod.rs @@ -22,6 +22,7 @@ pub mod expm1; pub mod factorial; pub mod floor; pub mod hex; +pub mod hypot; pub mod modulus; pub mod negative; pub mod pow; @@ -41,6 +42,7 @@ make_udf_function!(expm1::SparkExpm1, expm1); make_udf_function!(factorial::SparkFactorial, factorial); make_udf_function!(floor::SparkFloor, floor); make_udf_function!(hex::SparkHex, hex); +make_udf_function!(hypot::SparkHypot, hypot); make_udf_function!(modulus::SparkMod, modulus); make_udf_function!(modulus::SparkPmod, pmod); make_udf_function!(pow::SparkPow, pow); @@ -66,6 +68,7 @@ pub mod expr_fn { )); export_functions!((floor, "Returns floor of expr.", arg1)); export_functions!((hex, "Computes hex value of the given column.", arg1)); + export_functions!((hypot, "Returns sqrt(a^2 + b^2) without intermediate overflow or underflow.", arg1 arg2)); export_functions!((modulus, "Returns the remainder of division of the first argument by the second argument.", arg1 arg2)); export_functions!((pmod, "Returns the positive remainder of division of the first argument by the second argument.", arg1 arg2)); export_functions!(( @@ -107,6 +110,7 @@ pub fn functions() -> Vec> { factorial(), floor(), hex(), + hypot(), modulus(), pmod(), pow(), diff --git a/datafusion/sqllogictest/test_files/spark/math/hypot.slt b/datafusion/sqllogictest/test_files/spark/math/hypot.slt index 1349be0a95ee7..3d9d03da7a37e 100644 --- a/datafusion/sqllogictest/test_files/spark/math/hypot.slt +++ b/datafusion/sqllogictest/test_files/spark/math/hypot.slt @@ -21,7 +21,45 @@ # For more information, please see: # https://github.com/apache/datafusion/issues/15914 -## Original Query: SELECT hypot(3, 4); -## PySpark 3.5.5 Result: {'HYPOT(3, 4)': 5.0, 'typeof(HYPOT(3, 4))': 'double', 'typeof(3)': 'int', 'typeof(4)': 'int'} -#query -#SELECT hypot(3::int, 4::int); +# Scalar: classic Pythagorean triples (3-4-5, 5-12-13) +query R +SELECT hypot(3, 4); +---- +5 + +query R +SELECT hypot(5, 12); +---- +13 + +# Double inputs +query R +SELECT hypot(3.0::double, 4.0::double); +---- +5 + +# NULL if either argument is NULL (binary kernel null propagation) +query R +SELECT hypot(NULL::double, 4.0::double); +---- +NULL + +query R +SELECT hypot(3.0::double, NULL::double); +---- +NULL + +# Array path, including a NULL row +query R +SELECT hypot(a, b) FROM (VALUES (3.0::double, 4.0::double), (6.0::double, 8.0::double), (NULL::double, 1.0::double)) AS t(a, b); +---- +5 +10 +NULL + +# Overflow-safe: a naive sqrt(a*a + b*b) would overflow to Infinity here, +# but f64::hypot (like Spark's Math.hypot) stays finite. +query B +SELECT hypot(3e200::double, 4e200::double) < 'Infinity'::double; +---- +true From 3da5c126e1be69300daf3bdae34f661896e5747c Mon Sep 17 00:00:00 2001 From: KarpagamKarthikeyan Date: Fri, 24 Jul 2026 19:26:35 -0700 Subject: [PATCH 2/2] Refactored to make_scalar_function + a spark_hypot helper, more tests --- datafusion/spark/src/function/math/hypot.rs | 28 +++---- .../test_files/spark/math/hypot.slt | 76 ++++++++++++++++++- 2 files changed, 84 insertions(+), 20 deletions(-) diff --git a/datafusion/spark/src/function/math/hypot.rs b/datafusion/spark/src/function/math/hypot.rs index 2f88a3fd9559d..a1e30a7e4abe2 100644 --- a/datafusion/spark/src/function/math/hypot.rs +++ b/datafusion/spark/src/function/math/hypot.rs @@ -25,6 +25,7 @@ use datafusion_common::utils::take_function_args; use datafusion_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; +use datafusion_functions::utils::make_scalar_function; /// Spark-compatible `hypot` function. /// @@ -46,8 +47,7 @@ impl Default for SparkHypot { impl SparkHypot { pub fn new() -> Self { Self { - // Spark only defines hypot over doubles; `exact` makes coercion - // guarantee both inputs are Float64 before `invoke` runs. + // Spark only defines hypot over doubles signature: Signature::exact( vec![DataType::Float64, DataType::Float64], Volatility::Immutable, @@ -70,21 +70,15 @@ impl ScalarUDFImpl for SparkHypot { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let num_rows = args.number_rows; - let [x, y] = take_function_args(self.name(), args.args)?; - - // Broadcast scalars to arrays so one path covers every combination. - let x = x.to_array(num_rows)?; - let y = y.to_array(num_rows)?; - - // Safe: the `exact` signature guarantees Float64 inputs. - let x = x.as_primitive::(); - let y = y.as_primitive::(); + make_scalar_function(spark_hypot, vec![])(&args.args) + } +} - // `binary` applies the op element-wise and returns NULL when either - // input is NULL — matching Spark's null semantics. - let result: Float64Array = binary(x, y, |a, b| a.hypot(b))?; +fn spark_hypot(args: &[ArrayRef]) -> Result { + let [x, y] = take_function_args("hypot", args)?; - Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) - } + let x = x.as_primitive::(); + let y = y.as_primitive::(); + let result: Float64Array = binary(x, y, |a, b| a.hypot(b))?; + Ok(Arc::new(result)) } diff --git a/datafusion/sqllogictest/test_files/spark/math/hypot.slt b/datafusion/sqllogictest/test_files/spark/math/hypot.slt index 3d9d03da7a37e..564b34add8b9f 100644 --- a/datafusion/sqllogictest/test_files/spark/math/hypot.slt +++ b/datafusion/sqllogictest/test_files/spark/math/hypot.slt @@ -38,7 +38,7 @@ SELECT hypot(3.0::double, 4.0::double); ---- 5 -# NULL if either argument is NULL (binary kernel null propagation) +# NULL if either argument is NULL query R SELECT hypot(NULL::double, 4.0::double); ---- @@ -57,9 +57,79 @@ SELECT hypot(a, b) FROM (VALUES (3.0::double, 4.0::double), (6.0::double, 8.0::d 10 NULL -# Overflow-safe: a naive sqrt(a*a + b*b) would overflow to Infinity here, -# but f64::hypot (like Spark's Math.hypot) stays finite. +# Overflow-safe: naive sqrt(a*a + b*b) overflows to Infinity here; hypot stays finite (matches Spark's Math.hypot) query B SELECT hypot(3e200::double, 4e200::double) < 'Infinity'::double; ---- true + +# any infinite input yields +Infinity, even when the other is NaN +query R +SELECT hypot('Infinity'::double, 4.0::double); +---- +Infinity + +query R +SELECT hypot(4.0::double, '-Infinity'::double); +---- +Infinity + +query R +SELECT hypot('Infinity'::double, 'NaN'::double); +---- +Infinity + +# NaN propagates when neither input is infinite +query R +SELECT hypot('NaN'::double, 4.0::double); +---- +NaN + +# signed zeros +query RRR +SELECT hypot(0.0::double, 0.0::double), hypot(-0.0::double, 0.0::double), hypot(3.0::double, -0.0::double); +---- +0 0 3 + +# NULL propagates even when the other input is Infinity +query R +SELECT hypot(NULL::double, 'Infinity'::double); +---- +NULL + +# negative inputs yield the positive magnitude +query RR +SELECT hypot(-3.0::double, -4.0::double), hypot(-3.0::double, 4.0::double); +---- +5 5 + +# Underflow-safe: naive sqrt(a*a + b*b) underflows to 0 for tiny inputs; hypot stays nonzero (matches Spark's Math.hypot) +query B +SELECT hypot(3e-200::double, 4e-200::double) > 0; +---- +true + +# Array path with special values (normal, +Infinity, NaN, NULL) +query R +SELECT hypot(a, b) FROM (VALUES + (3.0::double, 4.0::double), + ('Infinity'::double, 1.0::double), + ('NaN'::double, 1.0::double), + (NULL::double, 1.0::double)) AS t(a, b); +---- +5 +Infinity +NaN +NULL + +# both inputs NaN -> NaN +query R +SELECT hypot('NaN'::double, 'NaN'::double); +---- +NaN + +# both inputs infinite -> +Infinity +query R +SELECT hypot('Infinity'::double, '-Infinity'::double); +---- +Infinity \ No newline at end of file