datafusion_spark/function/math/
hypot.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use arrow::array::{ArrayRef, AsArray, Float64Array};
21use arrow::compute::kernels::arity::binary;
22use arrow::datatypes::{DataType, Float64Type};
23use datafusion_common::Result;
24use datafusion_common::utils::take_function_args;
25use datafusion_expr::{
26    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
27};
28use datafusion_functions::utils::make_scalar_function;
29
30/// Spark-compatible `hypot` function.
31///
32/// <https://spark.apache.org/docs/latest/api/sql/index.html#hypot>
33///
34/// Returns `sqrt(expr1^2 + expr2^2)` computed without intermediate overflow or
35/// underflow, matching Spark's use of `java.lang.Math.hypot`.
36#[derive(Debug, PartialEq, Eq, Hash)]
37pub struct SparkHypot {
38    signature: Signature,
39}
40
41impl Default for SparkHypot {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl SparkHypot {
48    pub fn new() -> Self {
49        Self {
50            // Spark only defines hypot over doubles
51            signature: Signature::exact(
52                vec![DataType::Float64, DataType::Float64],
53                Volatility::Immutable,
54            ),
55        }
56    }
57}
58
59impl ScalarUDFImpl for SparkHypot {
60    fn name(&self) -> &str {
61        "hypot"
62    }
63
64    fn signature(&self) -> &Signature {
65        &self.signature
66    }
67
68    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
69        Ok(DataType::Float64)
70    }
71
72    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
73        make_scalar_function(spark_hypot, vec![])(&args.args)
74    }
75}
76
77fn spark_hypot(args: &[ArrayRef]) -> Result<ArrayRef> {
78    let [x, y] = take_function_args("hypot", args)?;
79
80    let x = x.as_primitive::<Float64Type>();
81    let y = y.as_primitive::<Float64Type>();
82    let result: Float64Array = binary(x, y, |a, b| a.hypot(b))?;
83    Ok(Arc::new(result))
84}
datafusion_spark/function/math/hypot.rs

datafusion_spark/function/math/
hypot.rs
