From 1978af26642ed3f73f06f91feabb6505452aa161 Mon Sep 17 00:00:00 2001 From: Shrivas Shankar Date: Tue, 16 Jun 2026 10:50:32 -0500 Subject: [PATCH] [C++][Python] Add "hypot" compute kernel Add a binary floating-point compute function "hypot" that computes the hypotenuse sqrt(x^2 + y^2) without undue overflow or underflow at intermediate stages, mirroring numpy.hypot and std::hypot. The kernel reuses the existing MakeArithmeticFunctionFloatingPoint factory (as atan2 does), registering float32 and float64 kernels and promoting integer/decimal inputs to float64 via DispatchBest. Also adds the Hypot() convenience function and declaration, a FunctionDoc, C++ tests (including an overflow-safety test that exercises inputs whose squares overflow float32), a compute.rst entry, and pyarrow expression coverage. --- cpp/src/arrow/compute/api_scalar.cc | 1 + cpp/src/arrow/compute/api_scalar.h | 9 ++++ .../compute/kernels/scalar_arithmetic.cc | 21 +++++++++ .../compute/kernels/scalar_arithmetic_test.cc | 44 +++++++++++++++++++ docs/source/cpp/compute.rst | 6 +++ python/pyarrow/tests/test_compute.py | 1 + 6 files changed, 82 insertions(+) diff --git a/cpp/src/arrow/compute/api_scalar.cc b/cpp/src/arrow/compute/api_scalar.cc index b43eca542f36..0aa8fd757a98 100644 --- a/cpp/src/arrow/compute/api_scalar.cc +++ b/cpp/src/arrow/compute/api_scalar.cc @@ -805,6 +805,7 @@ SCALAR_ARITHMETIC_BINARY(ShiftLeft, "shift_left", "shift_left_checked") SCALAR_ARITHMETIC_BINARY(ShiftRight, "shift_right", "shift_right_checked") SCALAR_ARITHMETIC_BINARY(Subtract, "subtract", "subtract_checked") SCALAR_EAGER_BINARY(Atan2, "atan2") +SCALAR_EAGER_BINARY(Hypot, "hypot") SCALAR_EAGER_UNARY(Floor, "floor") SCALAR_EAGER_UNARY(Ceil, "ceil") SCALAR_EAGER_UNARY(Trunc, "trunc") diff --git a/cpp/src/arrow/compute/api_scalar.h b/cpp/src/arrow/compute/api_scalar.h index 8b341e865a16..c4238b956c9c 100644 --- a/cpp/src/arrow/compute/api_scalar.h +++ b/cpp/src/arrow/compute/api_scalar.h @@ -806,6 +806,15 @@ Result Atan(const Datum& arg, ExecContext* ctx = NULLPTR); ARROW_EXPORT Result Atan2(const Datum& y, const Datum& x, ExecContext* ctx = NULLPTR); +/// \brief Compute the hypotenuse (Euclidean norm) of x and y, equivalent to +/// sqrt(x^2 + y^2), without undue overflow or underflow at intermediate stages. +/// \param[in] x The x-values to compute the hypotenuse for. +/// \param[in] y The y-values to compute the hypotenuse for. +/// \param[in] ctx the function execution context, optional +/// \return the elementwise hypotenuse of the values +ARROW_EXPORT +Result Hypot(const Datum& x, const Datum& y, ExecContext* ctx = NULLPTR); + /// \brief Compute the hyperbolic sine of the array values. /// \param[in] arg The values to compute the hyperbolic sine for. /// \param[in] ctx the function execution context, optional diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index f09e209e81df..de1aebcd12bd 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -368,6 +368,15 @@ struct Atan2 { } }; +struct Hypot { + template + static enable_if_floating_value Call(KernelContext*, Arg0 x, Arg1 y, Status*) { + static_assert(std::is_same::value, ""); + static_assert(std::is_same::value, ""); + return std::hypot(x, y); + } +}; + struct LogNatural { template static enable_if_floating_value Call(KernelContext*, Arg arg, Status*) { @@ -1346,6 +1355,14 @@ const FunctionDoc atan2_doc{"Compute the inverse tangent of y/x", ("The return value is in the range [-pi, pi]."), {"y", "x"}}; +const FunctionDoc hypot_doc{ + "Compute the hypotenuse (Euclidean norm) of x and y", + ("The result is equivalent to `sqrt(x^2 + y^2)`, but is computed without\n" + "undue overflow or underflow at intermediate stages of the computation.\n" + "If either x or y is +/-infinity, +infinity is returned, even if the\n" + "other argument is NaN."), + {"x", "y"}}; + const FunctionDoc atanh_doc{"Compute the inverse hyperbolic tangent", ("NaN is returned for input values x with \\|x\\| > 1.\n" "At x = +/- 1, returns +/- infinity.\n" @@ -1765,6 +1782,10 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { "sqrt_checked", sqrt_checked_doc); DCHECK_OK(registry->AddFunction(std::move(sqrt_checked))); + // ---------------------------------------------------------------------- + auto hypot = MakeArithmeticFunctionFloatingPoint("hypot", hypot_doc); + DCHECK_OK(registry->AddFunction(std::move(hypot))); + // ---------------------------------------------------------------------- auto sign = MakeUnaryArithmeticFunctionWithFixedIntOutType("sign", sign_doc); diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc index 9367ad2c89d1..9883ee62b7cc 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic_test.cc @@ -2807,6 +2807,50 @@ TYPED_TEST(TestBinaryArithmeticFloating, TrigAtan2) { -M_PI_2, 0, M_PI)); } +TYPED_TEST(TestBinaryArithmeticFloating, Hypot) { + SKIP_IF_HALF_FLOAT(); + + this->SetNansEqual(true); + auto hypot = [](const Datum& x, const Datum& y, ArithmeticOptions, ExecContext* ctx) { + return Hypot(x, y, ctx); + }; + this->AssertBinop(hypot, "[]", "[]", "[]"); + // Pythagorean triples; result is independent of the sign of either argument, + // and hypot(0, 0) == 0. + this->AssertBinop(hypot, "[3, -3, 5, -8, 0]", "[4, -4, -12, 15, 0]", + "[5, 5, 13, 17, 0]"); + // Null propagation. + this->AssertBinop(hypot, "[1, null, 0]", "[null, 1, 0]", "[null, null, 0]"); + // NaN propagates, unless the other argument is infinite (per C99/IEEE 754, + // hypot(+/-Inf, NaN) == +Inf). + this->AssertBinop(hypot, "[NaN, 1, NaN, Inf]", "[1, NaN, NaN, NaN]", + "[NaN, NaN, NaN, Inf]"); + // +/-infinity in either argument yields +infinity. + this->AssertBinop(hypot, "[Inf, -Inf, 3]", "[4, 0, -Inf]", "[Inf, Inf, Inf]"); +} + +// hypot avoids overflow/underflow at intermediate stages: for float32 the +// squares below overflow to +Inf, so a naive sqrt(x*x + y*y) would return Inf, +// while the kernel (like std::hypot) returns the correct finite result. +TEST(TestBinaryArithmetic, HypotOverflowSafety) { + std::vector xs = {3.0e30f, 5.0e37f, -2.0e30f}; + std::vector ys = {4.0e30f, 1.2e38f, 0.0f}; + ASSERT_TRUE(std::isinf(xs[0] * xs[0])); // the naive intermediate overflows + + std::vector expected_vals; + for (size_t i = 0; i < xs.size(); ++i) { + expected_vals.push_back(std::hypot(xs[i], ys[i])); + } + + std::shared_ptr x, y, expected; + ArrayFromVector(xs, &x); + ArrayFromVector(ys, &y); + ArrayFromVector(expected_vals, &expected); + + ASSERT_OK_AND_ASSIGN(Datum result, Hypot(x, y)); + AssertArraysEqual(*expected, *result.make_array(), /*verbose=*/true); +} + TYPED_TEST(TestUnaryArithmeticFloating, TrigAtanh) { SKIP_IF_HALF_FLOAT(); diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index e4092af70cde..1e067c52188d 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -512,6 +512,8 @@ Mixed time resolution temporal inputs will be cast to finest input resolution. +------------------+--------+-------------------------+-------------------------------+-------+ | expm1 | Unary | Numeric | Float32/Float64 | | +------------------+--------+-------------------------+-------------------------------+-------+ +| hypot | Binary | Numeric | Float32/Float64 | \(3) | ++------------------+--------+-------------------------+-------------------------------+-------+ | multiply | Binary | Numeric/Temporal | Numeric/Temporal | \(1) | +------------------+--------+-------------------------+-------------------------------+-------+ | multiply_checked | Binary | Numeric/Temporal | Numeric/Temporal | \(1) | @@ -560,6 +562,10 @@ Mixed time resolution temporal inputs will be cast to finest input resolution. values return NaN. Integral and decimal values return signedness as Int8 and floating-point values return it with the same type as the input values. +* \(3) Computes ``sqrt(x^2 + y^2)`` without undue overflow or underflow at + intermediate stages of the computation. If either argument is infinite, the + result is ``+Inf`` even if the other argument is NaN. + Bit-wise functions ~~~~~~~~~~~~~~~~~~ diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 478a0c3dc53c..73c30c7946a5 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -3932,6 +3932,7 @@ def create_sample_expressions(): pc.multiply(a, b), pc.power(a, a), pc.sqrt(a), pc.exp(b), pc.cos(b), pc.sin(b), pc.tan(b), pc.acos(b), pc.atan(b), pc.asin(b), pc.atan2(b, b), + pc.hypot(b, b), pc.sinh(a), pc.cosh(a), pc.tanh(a), pc.asinh(a), pc.acosh(b), pc.atanh(k), pc.abs(b), pc.sign(a), pc.bit_wise_not(a),