diff --git a/src/postgres/math_udfs.rs b/src/postgres/math_udfs.rs index 97bf072..b8d4d15 100644 --- a/src/postgres/math_udfs.rs +++ b/src/postgres/math_udfs.rs @@ -5,7 +5,7 @@ use datafusion::arrow::datatypes::DataType; use datafusion::common::DataFusionError; use datafusion::error::Result; -/// Inverse cosine, result in radians. +/// Inverse cosine, result in degrees. pub fn acosd(args: &[ArrayRef]) -> Result { let values = datafusion::common::cast::as_float64_array(&args[0])?; let mut float64array_builder = Float64Array::builder(args[0].len()); @@ -37,6 +37,38 @@ pub fn acosd(args: &[ArrayRef]) -> Result { Ok(Arc::new(float64array_builder.finish()) as ArrayRef) } +/// Inverse cosine, result in radians. +pub fn acos(args: &[ArrayRef]) -> Result { + 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 { + if value > 1.0 { + return Err(DataFusionError::Internal( + "input is out of range".to_string(), + )); + } + let result = value.acos(); + if result.fract() < 0.9 { + if result.fract() < 0.01 { + float64array_builder.append_value(result.floor()); + } else { + float64array_builder.append_value(result); + } + } else { + float64array_builder.append_value(result.ceil()); + } + Ok::<(), DataFusionError>(()) + } else { + float64array_builder.append_null(); + Ok::<(), DataFusionError>(()) + } + })?; + + Ok(Arc::new(float64array_builder.finish()) as ArrayRef) +} + /// Nearest integer greater than or equal to argument (same as ceil). pub fn ceiling(args: &[ArrayRef]) -> Result { let values = datafusion::common::cast::as_float64_array(&args[0])?; @@ -238,6 +270,63 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_acos() -> Result<()> { + let ctx = register_udfs_for_test()?; + let df = ctx.sql("select acos(0.5) as col_result").await?; + + let batches = df.clone().collect().await?; + + let expected: Vec<&str> = r#" ++--------------------+ +| col_result | ++--------------------+ +| 1.0471975511965976 | ++--------------------+"# + .split('\n') + .filter_map(|input| { + if input.is_empty() { + None + } else { + Some(input.trim()) + } + }) + .collect(); + assert_batches_sorted_eq!(expected, &batches); + + let df = ctx.sql("select acos(0.4) as col_result").await?; + + let batches = df.clone().collect().await?; + + let expected: Vec<&str> = r#" ++--------------------+ +| col_result | ++--------------------+ +| 1.1592794807274085 | ++--------------------+"# + .split('\n') + .filter_map(|input| { + if input.is_empty() { + None + } else { + Some(input.trim()) + } + }) + .collect(); + assert_batches_sorted_eq!(expected, &batches); + + let df = ctx.sql("select acos(1.4) as col_result").await?; + + let result = df.clone().collect().await; + assert!(result + .err() + .unwrap() + .to_string() + .contains("input is out of range")); + + Ok(()) + } + #[tokio::test] async fn test_ceiling() -> Result<()> { let ctx = register_udfs_for_test()?; diff --git a/src/postgres/mod.rs b/src/postgres/mod.rs index d085dbb..4c2113b 100644 --- a/src/postgres/mod.rs +++ b/src/postgres/mod.rs @@ -2,7 +2,7 @@ use std::sync::Arc; -use crate::postgres::math_udfs::{acosd, ceiling, div, erf, erfc}; +use crate::postgres::math_udfs::{acos, acosd, ceiling, div, erf, erfc}; use crate::postgres::network_udfs::{ broadcast, family, host, hostmask, inet_merge, inet_same_family, masklen, netmask, network, set_masklen, @@ -24,6 +24,7 @@ pub fn register_postgres_udfs(ctx: &SessionContext) -> Result<()> { fn register_math_udfs(ctx: &SessionContext) -> Result<()> { register_acosd(ctx); + register_acos(ctx); register_ceiling(ctx); register_erf(ctx); register_erfc(ctx); @@ -44,6 +45,19 @@ fn register_acosd(ctx: &SessionContext) { ctx.register_udf(acosd_udf); } +fn register_acos(ctx: &SessionContext) { + let acos_udf = make_scalar_function(acos); + let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(Float64))); + let acos_udf = ScalarUDF::new( + "acos", + &Signature::uniform(1, vec![Float64], Volatility::Immutable), + &return_type, + &acos_udf, + ); + + ctx.register_udf(acos_udf); +} + fn register_ceiling(ctx: &SessionContext) { let ceiling_udf = make_scalar_function(ceiling); let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(Float64))); diff --git a/supports/postgres.md b/supports/postgres.md index d6aff06..750e597 100644 --- a/supports/postgres.md +++ b/supports/postgres.md @@ -37,7 +37,7 @@ https://www.postgresql.org/docs/16/functions-math.html | ✅ | acosd ( double precision ) → double precision | Inverse cosine, result in degrees | acosd(0.5) → 60 | | ❓ | asind ( double precision ) → double precision | Inverse sine, result in degrees | asind(0.5) → 30 | | ❓ | atand ( double precision ) → double precision | Inverse tangent, result in degrees | atand(1) → 45 | -| ❓ | cosd ( double precision ) → double precision | Cosine, argument in degrees | cosd(60) → 0.5 | +| ✅ | cosd ( double precision ) → double precision | Cosine, argument in degrees | cosd(60) → 0.5 | | ❓ | cotd ( double precision ) → double precision | Cotangent, argument in degrees | cotd(45) → 1 | | ❓ | sind ( double precision ) → double precision | Sine, argument in degrees | sind(30) → 0.5 | | ❓ | tand ( double precision ) → double precision | Tangent, argument in degrees | tand(45) → 1 |