Skip to content

Commit

Permalink
Added postgres atand
Browse files Browse the repository at this point in the history
  • Loading branch information
dadepo committed Dec 30, 2023
1 parent a12c4b2 commit ada6bbe
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 2 deletions.
76 changes: 76 additions & 0 deletions src/postgres/math_udfs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,34 @@ pub fn sind(args: &[ArrayRef]) -> Result<ArrayRef> {
Ok(Arc::new(float64array_builder.finish()) as ArrayRef)
}


/// Inverse tangent, result in degrees.
pub fn atand(args: &[ArrayRef]) -> Result<ArrayRef> {
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();
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<ArrayRef> {
let values = datafusion::common::cast::as_float64_array(&args[0])?;
Expand Down Expand Up @@ -452,6 +480,54 @@ mod tests {
Ok(())
}

#[tokio::test]
async fn test_atand() -> Result<()> {
let ctx = register_udfs_for_test()?;
let df = ctx.sql("select atand(0.5) as col_result").await?;

let batches = df.clone().collect().await?;

let expected: Vec<&str> = r#"
+--------------------+
| col_result |
+--------------------+
| 26.565051177077994 |
+--------------------+"#
.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 atand(1) as col_result").await?;

let batches = df.clone().collect().await?;

let expected: Vec<&str> = r#"
+------------+
| col_result |
+------------+
| 45.0 |
+------------+"#
.split('\n')
.filter_map(|input| {
if input.is_empty() {
None
} else {
Some(input.trim())
}
})
.collect();
assert_batches_sorted_eq!(expected, &batches);

Ok(())
}

#[tokio::test]
async fn test_ceiling() -> Result<()> {
let ctx = register_udfs_for_test()?;
Expand Down
16 changes: 15 additions & 1 deletion src/postgres/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

use std::sync::Arc;

use crate::postgres::math_udfs::{acosd, asind, ceiling, cosd, div, erf, erfc, sind};
use crate::postgres::math_udfs::{acosd, asind, atand, ceiling, cosd, div, erf, erfc, sind};
use crate::postgres::network_udfs::{
broadcast, family, host, hostmask, inet_merge, inet_same_family, masklen, netmask, network,
set_masklen,
Expand All @@ -27,6 +27,7 @@ fn register_math_udfs(ctx: &SessionContext) -> Result<()> {
register_cosd(ctx);
register_asind(ctx);
register_sind(ctx);
register_atand(ctx);
register_ceiling(ctx);
register_erf(ctx);
register_erfc(ctx);
Expand Down Expand Up @@ -86,6 +87,19 @@ fn register_cosd(ctx: &SessionContext) {
ctx.register_udf(cosd_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_ceiling(ctx: &SessionContext) {
let ceiling_udf = make_scalar_function(ceiling);
let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(Float64)));
Expand Down
2 changes: 1 addition & 1 deletion supports/postgres.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ https://www.postgresql.org/docs/16/functions-math.html
|| random_normal ( [ mean double precision [, stddev double precision ]] ) → double precision | Returns a random value from the normal distribution with the given parameters; mean defaults to 0.0 and stddev defaults to 1.0 | random_normal(0.0, 1.0) → 0.051285419 |
|| 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 |
| | 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 |
|| 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 |
Expand Down

0 comments on commit ada6bbe

Please sign in to comment.