Added new functions for outputting timevectors to templates

timescale · Nov 28, 2022 · 79fba01 · 79fba01
1 parent 0eca871
commit 79fba01
Show file tree

Hide file tree

Showing 3 changed files with 229 additions and 18 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/extension/Cargo.toml b/extension/Cargo.toml
@@ -43,6 +43,7 @@ rand = { version = "0.8.3", features = ["getrandom", "small_rng"] }
 rand_distr = "0.4.0"
 rand_chacha = "0.3.0"
 ron="0.6.0"
+tera = { version = "1.17.0", default-features = false }
 twofloat = { version = "0.6.0", features = ["serde"] }
 num-traits = "0.2.15"
 

diff --git a/extension/src/time_vector.rs b/extension/src/time_vector.rs
@@ -1,6 +1,10 @@
 #![allow(clippy::identity_op)] // clippy gets confused by pg_type! enums
 
+use crate::pg_sys::timestamptz_to_str;
+use core::str::Utf8Error;
 use pgx::{iter::TableIterator, *};
+use std::ffi::CStr;
+use tera::{Context, Tera};
 
 use crate::{
     aggregate_utils::in_aggregate_context,
@@ -115,6 +119,48 @@ pub fn unnest<'a>(
     )
 }
 
+/// Util function to convert from *const ::std::os::raw::c_char to String
+/// TimestampTz -> *const c_char -> &CStr -> &str -> String
+pub fn timestamptz_to_string(time: pg_sys::TimestampTz) -> Result<String, Utf8Error> {
+    let char_ptr = unsafe { timestamptz_to_str(time) };
+    let c_str = unsafe { CStr::from_ptr(char_ptr) };
+    c_str.to_str().map(|s| s.to_owned())
+}
+
+#[pg_extern(immutable, schema = "toolkit_experimental", parallel_safe)]
+pub fn format_timevector<'a>(series: Timevector_TSTZ_F64<'a>, format_string: String) -> String {
+    let mut context = Context::new();
+    let mut times: Vec<String> = Vec::new();
+    let mut values: Vec<String> = Vec::new();
+    if series.has_nulls() {
+        for (i, point) in series.iter().enumerate() {
+            times.push(timestamptz_to_string(point.ts).unwrap());
+            if series.is_null_val(i) {
+                values.push("null".to_string())
+            } else {
+                match point.val.to_string().as_ref() {
+                    "NaN" | "Inf" => panic!(),
+                    x => values.push(x.to_string()),
+                }
+            }
+        }
+    } else {
+        // optimized path if series does not have any nulls
+        for point in series {
+            times.push(timestamptz_to_string(point.ts).unwrap());
+            values.push(point.val.to_string());
+        }
+    }
+    context.insert("TIMES", &times);
+    context.insert("VALUES", &values);
+    match format_string.as_ref() {
+	"plotly" =>     Tera::one_off("{\"times\": {{ TIMES | json_encode(pretty=true) | safe  }}, \"vals\": {{ VALUES | json_encode(pretty=true) | safe }}}", &context, true).expect("Failed to create plotly template"),
+	"paired" =>    {
+	    Tera::one_off( "[{% for x in TIMES %}{\"time\": \"{{ x }}\", \"val\": {{ VALUES[loop.index0] }}}, {% endfor %}]", &context, true).expect("Failed to create paired template")}, 
+	format_string =>    Tera::one_off(&format_string, &context, true).expect("Failed to create template with Tera")
+    }
+}
+
 #[pg_operator(immutable, parallel_safe)]
 #[opname(->)]
 pub fn arrow_timevector_unnest<'a>(
@@ -424,6 +470,52 @@ mod tests {
         })
     }
 
+    #[pg_test]
+    pub fn test_format_timevector() {
+        Spi::execute(|client| {
+            client.select("SET timezone TO 'UTC'", None, None);
+            client.select(
+                "CREATE TABLE data(time TIMESTAMPTZ, value DOUBLE PRECISION)",
+                None,
+                None,
+            );
+            client.select(
+                r#"INSERT INTO data VALUES
+                    ('2020-1-1', 30.0),
+                    ('2020-1-2', 45.0),
+                    ('2020-1-3', NULL),
+                    ('2020-1-4', 55.5),
+                    ('2020-1-5', 10.0)"#,
+                None,
+                None,
+            );
+
+            let test_plotly_template = client.select(
+                // "SELECT toolkit_experimental.format_timevector(timevector(time, value),'{\"times\": {{ TIMES }}, \"vals\": {{ VALUES }}}') FROM data",
+                "SELECT toolkit_experimental.format_timevector(timevector(time, value),'plotly') FROM data",
+                None,
+                None,
+            ).first()
+                .get_one::<String>()
+                .unwrap();
+
+            assert_eq!(test_plotly_template,"{\"times\": [\n  \"2020-01-01 00:00:00+00\",\n  \"2020-01-02 00:00:00+00\",\n  \"2020-01-03 00:00:00+00\",\n  \"2020-01-04 00:00:00+00\",\n  \"2020-01-05 00:00:00+00\"\n], \"vals\": [\n  \"30\",\n  \"45\",\n  \"null\",\n  \"55.5\",\n  \"10\"\n]}"
+		     );
+
+            let test_paired_template = client.select(
+                "SELECT toolkit_experimental.format_timevector(timevector(time, value),'paired') FROM data",
+                None,
+                None,
+            ).first()
+                .get_one::<String>()
+                .unwrap();
+
+            assert_eq!(
+                test_paired_template,
+        	"[{\"time\": \"2020-01-01 00:00:00+00\", \"val\": 30}, {\"time\": \"2020-01-02 00:00:00+00\", \"val\": 45}, {\"time\": \"2020-01-03 00:00:00+00\", \"val\": null}, {\"time\": \"2020-01-04 00:00:00+00\", \"val\": 55.5}, {\"time\": \"2020-01-05 00:00:00+00\", \"val\": 10}, ]"
+            );
+        })
+    }
     #[pg_test]
     pub fn timevector_io() {
         Spi::execute(|client| {