Switch to first-class aggregate support in pgx

Note this is a WIP, and removes backwards-compatibility functions.

Author: JamesGuthrie

Cargo.lock

@@ -1,31 +1,8 @@
 use crate::aggregates::{Microseconds, Milliseconds, STALE_NAN, USECS_PER_MS, USECS_PER_SEC};
-use crate::palloc::{Inner, InternalAsValue};
 use pgx::*;
 use serde::{Deserialize, Serialize};
 use std::collections::VecDeque;
 
-#[pg_extern(immutable, parallel_safe)]
-pub fn prom_extrapolate_final(state: Internal) -> Option<Vec<Option<f64>>> {
-    prom_extrapolate_final_inner(unsafe { state.to_inner() })
-}
-pub fn prom_extrapolate_final_inner(
-    state: Option<Inner<GapfillDeltaTransition>>,
-) -> Option<Vec<Option<f64>>> {
-    state.map(|mut s| s.as_vec())
-}
-
-/// Backwards compatibility
-#[no_mangle]
-pub extern "C" fn pg_finfo_gapfill_delta_final() -> &'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_final(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
-    prom_extrapolate_final_wrapper(fcinfo)
-}
-
 #[derive(Serialize, Deserialize, PostgresType, Debug)]
 pub struct GapfillDeltaTransition {
     window: VecDeque<(pg_sys::TimestampTz, f64)>,

src/aggregate_utils.rs

@@ -2,122 +2,56 @@ use pgx::*;
 
 use pgx::error;
 
-use crate::aggregate_utils::in_aggregate_context;
 use crate::aggregates::{GapfillDeltaTransition, Milliseconds};
-use crate::palloc::{Inner, InternalAsValue, ToInternal};
+
+#[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: pg_sys::TimestampTz,
-    greatest_time: pg_sys::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: pg_sys::TimestampTz,
-    sample_value: f64,
-    fc: pg_sys::FunctionCallInfo,
-) -> Internal {
-    prom_delta_transition_inner(
-        unsafe { state.to_inner() },
-        lowest_time,
-        greatest_time,
-        step_size,
-        range,
-        sample_time,
-        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 = (
+        name!(lowest_time, pg_sys::TimestampTz),
+        name!(greatest_time, pg_sys::TimestampTz),
+        name!(step_size, Milliseconds),
+        name!(range, Milliseconds),
+        name!(sample_time, pg_sys::TimestampTz),
+        name!(sample_value, 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,
+        _: pg_sys::FunctionCallInfo,
+    ) -> Self::State {
+        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, _: pg_sys::FunctionCallInfo) -> 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 {
@@ -164,7 +98,7 @@ mod tests {
         Spi::get_one::<Vec<f64>>(&*prepare_query("'2000-01-02 15:00:00 UTC'", "NULL"));
     }
 
-    #[pg_test(error = "input time less than lowest time")]
+    #[pg_test(error = "input time 631292400000000 not in bounds [140400000000, 143100000000]")]
     fn test_prom_delta_with_input_time_less_than_lowest_time_fails() {
         setup();
         Spi::get_one::<Vec<f64>>(&*prepare_query(