WIP: port remaining aggregates to new aggregate framework

Note: This depends on some changes to the aggregate implementation

Author: JamesGuthrie

Cargo.lock

@@ -36,6 +36,6 @@ serde_json = { version = "1.0.70", optional = true }
 pgx-tests = "0.2.5"
 
 [patch.crates-io]
-pgx = { git = "https://github.com/ZomboDB/pgx", branch = "aggregate" }
-pgx-macros = { git = "https://github.com/ZomboDB/pgx", branch = "aggregate" }
-pgx-tests = { git = "https://github.com/ZomboDB/pgx", branch = "aggregate" }
+pgx = { path = "../pgx/pgx" }
+pgx-macros = { path = "../pgx/pgx-macros" }
+pgx-tests = { path = "../pgx/pgx-tests" }

Cargo.toml

@@ -7,118 +7,64 @@ use crate::aggregates::{GapfillDeltaTransition, Milliseconds};
 use crate::palloc::{Inner, InternalAsValue, ToInternal};
 use crate::raw::TimestampTz;
 
+#[allow(non_camel_case_types)]
+pub struct prom_delta;
+
 // prom divides time into sliding windows of fixed size, e.g.
 // |  5 seconds  |  5 seconds  |  5 seconds  |  5 seconds  |  5 seconds  |
 // we take the first and last values in that bucket and uses `last-first` as the
 // value for that bucket.
 //  | a b c d e | f g h i | j   k |   m    |
 //  |   e - a   |  i - f  | k - j | <null> |
-#[allow(clippy::too_many_arguments)]
-#[pg_extern(immutable, parallel_safe)]
-pub fn prom_delta_transition(
-    state: Internal,
-    lowest_time: TimestampTz,
-    greatest_time: TimestampTz,
-    step_size: Milliseconds, // `prev_now - step_size` is where the next window starts
-    range: Milliseconds,     // the size of a window to delta over
-    sample_time: TimestampTz,
-    sample_value: f64,
-    fc: pg_sys::FunctionCallInfo,
-) -> Internal {
-    prom_delta_transition_inner(
-        unsafe { state.to_inner() },
-        lowest_time.into(),
-        greatest_time.into(),
-        step_size,
-        range,
-        sample_time.into(),
-        sample_value,
-        fc,
-    )
-    .internal()
-}
-
-#[allow(clippy::too_many_arguments)]
-fn prom_delta_transition_inner(
-    state: Option<Inner<GapfillDeltaTransition>>,
-    lowest_time: pg_sys::TimestampTz,
-    greatest_time: pg_sys::TimestampTz,
-    step_size: Milliseconds, // `prev_now - step` is where the next window starts
-    range: Milliseconds,     // the size of a window to delta over
-    sample_time: pg_sys::TimestampTz,
-    sample_value: f64,
-    fc: pg_sys::FunctionCallInfo,
-) -> Option<Inner<GapfillDeltaTransition>> {
-    unsafe {
-        in_aggregate_context(fc, || {
-            if sample_time < lowest_time || sample_time > greatest_time {
-                error!("input time less than lowest time")
-            }
-
-            let mut state = state.unwrap_or_else(|| {
-                let state: Inner<_> = GapfillDeltaTransition::new(
-                    lowest_time,
-                    greatest_time,
-                    range,
-                    step_size,
-                    false,
-                    false,
-                )
-                .into();
-                state
-            });
-
-            state.add_data_point(sample_time, sample_value);
-
-            Some(state)
-        })
+#[pg_aggregate]
+impl Aggregate for prom_delta {
+    type State = Option<GapfillDeltaTransition>;
+    type Args = (
+        TimestampTz,
+        TimestampTz,
+        Milliseconds,
+        Milliseconds,
+        TimestampTz,
+        f64,
+    );
+    type Finalize = Option<Vec<Option<f64>>>;
+
+    fn state(
+        state: Self::State,
+        (lowest_time, greatest_time, step_size, range, sample_time, sample_value): Self::Args,
+    ) -> Self::State {
+        let sample_time: pg_sys::TimestampTz = sample_time.into();
+        let lowest_time: pg_sys::TimestampTz = lowest_time.into();
+        let greatest_time: pg_sys::TimestampTz = greatest_time.into();
+
+        if sample_time < lowest_time || sample_time > greatest_time {
+            error!(format!(
+                "input time {} not in bounds [{}, {}]",
+                sample_time, lowest_time, greatest_time
+            ))
+        }
+
+        let mut state = state.unwrap_or_else(|| {
+            GapfillDeltaTransition::new(
+                lowest_time,
+                greatest_time,
+                range,
+                step_size,
+                false,
+                false,
+            )
+        });
+
+        state.add_data_point(sample_time, sample_value);
+
+        Some(state)
     }
-}
-
-/// Backwards compatibility
-#[no_mangle]
-pub extern "C" fn pg_finfo_gapfill_delta_transition() -> &'static pg_sys::Pg_finfo_record {
-    const V1_API: pg_sys::Pg_finfo_record = pg_sys::Pg_finfo_record { api_version: 1 };
-    &V1_API
-}
 
-#[no_mangle]
-unsafe extern "C" fn gapfill_delta_transition(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
-    prom_delta_transition_wrapper(fcinfo)
-}
-
-#[no_mangle]
-pub extern "C" fn pg_finfo_prom_delta_final_wrapper() -> &'static pg_sys::Pg_finfo_record {
-    const V1_API: pg_sys::Pg_finfo_record = pg_sys::Pg_finfo_record { api_version: 1 };
-    &V1_API
-}
-
-#[no_mangle]
-unsafe extern "C" fn prom_delta_final_wrapper(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
-    super::gapfill_delta::prom_extrapolate_final_wrapper(fcinfo)
+    fn finalize(current: Self::State) -> Self::Finalize {
+        current.map(|mut s| s.as_vec())
+    }
 }
 
-// implementation of prometheus delta function
-// for proper behavior the input must be ORDER BY sample_time
-extension_sql!(
-    r#"
-CREATE AGGREGATE @[email protected]_delta(
-    lowest_time TIMESTAMPTZ,
-    greatest_time TIMESTAMPTZ,
-    step_size BIGINT,
-    range BIGINT,
-    sample_time TIMESTAMPTZ,
-    sample_value DOUBLE PRECISION)
-(
-    sfunc=@[email protected]_delta_transition,
-    stype=internal,
-    finalfunc=@[email protected]_extrapolate_final
-);
-"#,
-    name = "create_prom_delta_aggregate",
-    requires = [prom_delta_transition, prom_extrapolate_final]
-);
-
 #[cfg(any(test, feature = "pg_test"))]
 #[pg_schema]
 mod tests {

src/aggregates/prom_delta.rs

@@ -26,11 +26,14 @@ impl Aggregate for prom_increase {
         let sample_time: pg_sys::TimestampTz = sample_time.into();
         let lowest_time: pg_sys::TimestampTz = lowest_time.into();
         let greatest_time: pg_sys::TimestampTz = greatest_time.into();
+
         if sample_time < lowest_time || sample_time > greatest_time {
-            error!("input time less than lowest time")
+            error!(format!(
+                "input time {} not in bounds [{}, {}]",
+                sample_time, lowest_time, greatest_time
+            ))
         }
 
-        // TODO: Missing in_aggregate_context magic
         let mut state = state.unwrap_or_else(|| {
             GapfillDeltaTransition::new(lowest_time, greatest_time, range, step_size, true, false)
         });