Module pgrx::aggregate

source ·
Expand description

Aggregate support.

Most items of this trait map directly to a CREATE AGGREGATE functionality.

Aggregates are created by implementing Aggregate for a type and decorating the implementation with #[pg_aggregate].

Definition of the aggregate is done via settings in the type’s Aggregate implementation. While the trait itself has several items, only a few are required, the macro will fill in the others with unused stubs.

§Minimal Example

use pgrx::prelude::*;
use serde::{Serialize, Deserialize};

// pgrx::pg_module_magic!(); // Uncomment this outside of docs!

#[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
pub struct DemoSum {
    count: i32,
}

#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = i32;
    fn state(
        mut current: Self::State,
        arg: Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        current.count += arg;
        current
    }
}

This creates SQL like so:

-- src/lib.rs:11
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
    integer /* i32 */
)
(
    SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
    STYPE = DemoSum, /* aggregate::DemoSum */
    INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);

Example of usage:

aggregate=# CREATE TABLE demo_table (value INTEGER);
CREATE TABLE
aggregate=# INSERT INTO demo_table (value) VALUES (1), (2), (3);
INSERT 0 3
aggregate=# SELECT DemoSum(value) FROM demo_table;
    demosum
-------------
 {"count":6}
(1 row)

§Multiple Arguments

Sometimes aggregates need to handle multiple arguments. The Aggregate::Args associated type can be a tuple:

#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = (i32, i32);
    fn state(
        mut current: Self::State,
        (arg1, arg2): Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        current.count += arg1;
        current.count += arg2;
        current
    }
}

Creates:

-- src/lib.rs:11
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
    integer, /* i32 */
    integer /* i32 */
)
(
    SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
    STYPE = DemoSum, /* aggregate::DemoSum */
    INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);

§Named Arguments

The name!(ident, Type) macro can be used to set the name of an argument:

impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = (
        i32,
        name!(extra, i32),
    );
    fn state(
        mut current: Self::State,
        (arg1, extra): Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        todo!()
    }
}

Creates:

-- src/lib.rs:11
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
    integer, /* i32 */
    "extra" integer /* i32 */
)
(
    SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
    STYPE = DemoSum, /* aggregate::DemoSum */
    INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);

§Function attributes

Functions inside the impl may use the #[pgrx] attribute. It accepts the same parameters as [#[pg_extern]][macro@pgrx-macros::pg_extern].

#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = i32;
    #[pgrx(parallel_safe, immutable)]
    fn state(
        mut current: Self::State,
        arg: Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        todo!()
    }
}

Generates:

-- src/lib.rs:11
-- aggregate::demo_sum_state
CREATE FUNCTION "demo_sum_state"(
    "this" DemoSum, /* aggregate::DemoSum */
    "arg_one" integer /* i32 */
) RETURNS DemoSum /* aggregate::DemoSum */
PARALLEL SAFE IMMUTABLE STRICT
LANGUAGE c /* Rust */
AS 'MODULE_PATHNAME', 'demo_sum_state_wrapper';

§Non-Self State

Sometimes it’s useful to have aggregates share state, or use some other type for state.

#[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
pub struct DemoSumState {
    count: i32,
}

pub struct DemoSum;

#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = i32;
    type State = DemoSumState;
    fn state(
        mut current: Self::State,
        arg: Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        todo!()
    }
}

Creates:

-- src/lib.rs:13
-- aggregate::demo_sum_state
CREATE FUNCTION "demo_sum_state"(
    "this" DemoSumState, /* aggregate::DemoSumState */
    "arg_one" integer /* i32 */
) RETURNS DemoSumState /* aggregate::DemoSumState */
STRICT
LANGUAGE c /* Rust */
AS 'MODULE_PATHNAME', 'demo_sum_state_wrapper';

-- src/lib.rs:13
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
    integer /* i32 */
)
(
    SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
    STYPE = DemoSumState, /* aggregate::DemoSumState */
    INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);

Enums§

Traits§