Added postgres tand

dadepo · Jan 6, 2024 · 219a1cf · 219a1cf
1 parent 02579c2
commit 219a1cf
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 2 deletions.
diff --git a/src/postgres/math_udfs.rs b/src/postgres/math_udfs.rs
@@ -157,6 +157,25 @@ pub fn atand(args: &[ArrayRef]) -> Result<ArrayRef> {
     Ok(Arc::new(float64array_builder.finish()) as ArrayRef)
 }
 
+/// Tangent, argument in degrees.
+pub fn tand(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.to_radians().tan();
+            float64array_builder.append_value(result);
+            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])?;
@@ -598,6 +617,54 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_tand() -> Result<()> {
+        let ctx = register_udfs_for_test()?;
+        let df = ctx.sql("select tand(45) as col_result").await?;
+
+        let batches = df.clone().collect().await?;
+
+        let expected: Vec<&str> = r#"
++--------------------+
+| col_result         |
++--------------------+
+| 0.9999999999999999 |
++--------------------+"#
+            .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 tand(0.4) as col_result").await?;
+
+        let batches = df.clone().collect().await?;
+
+        let expected: Vec<&str> = r#"
++----------------------+
+| col_result           |
++----------------------+
+| 0.006981430430496479 |
++----------------------+"#
+            .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()?;

diff --git a/src/postgres/mod.rs b/src/postgres/mod.rs
@@ -2,7 +2,9 @@
 
 use std::sync::Arc;
 
-use crate::postgres::math_udfs::{acosd, asind, atand, ceiling, cosd, cotd, div, erf, erfc, sind};
+use crate::postgres::math_udfs::{
+    acosd, asind, atand, ceiling, cosd, cotd, div, erf, erfc, sind, tand,
+};
 use crate::postgres::network_udfs::{
     broadcast, family, host, hostmask, inet_merge, inet_same_family, masklen, netmask, network,
     set_masklen,
@@ -29,6 +31,7 @@ fn register_math_udfs(ctx: &SessionContext) -> Result<()> {
     register_asind(ctx);
     register_sind(ctx);
     register_atand(ctx);
+    register_tand(ctx);
     register_ceiling(ctx);
     register_erf(ctx);
     register_erfc(ctx);
@@ -114,6 +117,19 @@ fn register_atand(ctx: &SessionContext) {
     ctx.register_udf(atand_udf);
 }
 
+fn register_tand(ctx: &SessionContext) {
+    let tand_udf = make_scalar_function(tand);
+    let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(Float64)));
+    let tand_udf = ScalarUDF::new(
+        "tand",
+        &Signature::uniform(1, vec![Float64], Volatility::Immutable),
+        &return_type,
+        &tand_udf,
+    );
+
+    ctx.register_udf(tand_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
@@ -40,4 +40,4 @@ https://www.postgresql.org/docs/16/functions-math.html
 | ✅           | 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                            | 
+| ✅           | tand ( double precision ) → double precision                                               | Tangent, argument in degrees                                                                                                                                                                  | tand(45) → 1                            |