diff --git a/src/postgres/math_udfs.rs b/src/postgres/math_udfs.rs index f8e632c..92f7974 100644 --- a/src/postgres/math_udfs.rs +++ b/src/postgres/math_udfs.rs @@ -119,22 +119,56 @@ pub fn sind(args: &[ArrayRef]) -> Result { } /// Inverse tangent, result in degrees. -pub fn atand(args: &[ArrayRef]) -> Result { - let values = datafusion::common::cast::as_float64_array(&args[0])?; - let mut float64array_builder = Float64Array::builder(args[0].len()); +#[derive(Debug)] +pub struct Atand { + signature: Signature, +} - values.iter().try_for_each(|value| { - if let Some(value) = value { - let result = value.atan().to_degrees(); - float64array_builder.append_value(result); - Ok::<(), DataFusionError>(()) - } else { - float64array_builder.append_null(); - Ok::<(), DataFusionError>(()) +impl Atand { + pub fn new() -> Self { + Self { + signature: Signature::uniform(1, vec![Float64], Volatility::Immutable), } - })?; + } +} - Ok(Arc::new(float64array_builder.finish()) as ArrayRef) +impl ScalarUDFImpl for Atand { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn name(&self) -> &str { + "atand" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(Float64) + } + + fn invoke(&self, args: &[ColumnarValue]) -> Result { + let args = ColumnarValue::values_to_arrays(args)?; + let values = datafusion::common::cast::as_float64_array(&args[0])?; + let mut float64array_builder = Float64Array::builder(args[0].len()); + + values.iter().try_for_each(|value| { + if let Some(value) = value { + let result = value.atan().to_degrees(); + float64array_builder.append_value(result); + Ok::<(), DataFusionError>(()) + } else { + float64array_builder.append_null(); + Ok::<(), DataFusionError>(()) + } + })?; + + Ok(ColumnarValue::Array( + Arc::new(float64array_builder.finish()) as ArrayRef, + )) + } } /// Tangent, argument in degrees. diff --git a/src/postgres/mod.rs b/src/postgres/mod.rs index 5bd6f9b..bd227a1 100644 --- a/src/postgres/mod.rs +++ b/src/postgres/mod.rs @@ -10,7 +10,7 @@ use datafusion::physical_expr::functions::make_scalar_function; use datafusion::prelude::SessionContext; use crate::postgres::math_udfs::{ - acosd, asind, atand, cosd, cotd, sind, Ceiling, Div, Erf, Erfc, RandomNormal, Tand, + acosd, asind, cosd, cotd, sind, Atand, Ceiling, Div, Erf, Erfc, RandomNormal, Tand, }; use crate::postgres::network_udfs::{ broadcast, family, host, hostmask, inet_merge, inet_same_family, masklen, netmask, network, @@ -32,7 +32,7 @@ fn register_math_udfs(ctx: &SessionContext) -> Result<()> { register_cotd(ctx); register_asind(ctx); register_sind(ctx); - register_atand(ctx); + ctx.register_udf(ScalarUDF::from(Atand::new())); ctx.register_udf(ScalarUDF::from(Tand::new())); ctx.register_udf(ScalarUDF::from(Ceiling::new())); ctx.register_udf(ScalarUDF::from(Div::new())); @@ -107,19 +107,6 @@ fn register_cotd(ctx: &SessionContext) { ctx.register_udf(cotd_udf); } -fn register_atand(ctx: &SessionContext) { - let atand_udf = make_scalar_function(atand); - let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(Float64))); - let atand_udf = ScalarUDF::new( - "atand", - &Signature::uniform(1, vec![Float64], Volatility::Immutable), - &return_type, - &atand_udf, - ); - - ctx.register_udf(atand_udf); -} - fn register_network_udfs(ctx: &SessionContext) -> Result<()> { register_broadcast(ctx); register_family(ctx);