Rust develops postgres extensions

Posted by CapEsiouS on Tue, 30 Nov 2021 10:21:31 +0100

preface

Rust language is a general system level programming language. It has no GC and is famous for ensuring memory security, concurrency security and high performance. It has been privately developed by Graydon Hoare since 2008 and sponsored by Mozilla in 2009. Version 0.1.0 was released for the first time in 2010 for the development of Servo engine. Version 1.0 was released on May 15, 2015. Since its release, up to 2021, after six years of development, rust has risen steadily and has gradually become mature and stable. From 2016 to 2021, rust has become the most popular language on the StackOverflow language list for five consecutive years [1]. On February 9, 2021, rust foundation was announced to be established. Leading technology giants such as Huawei, AWS, Google, Microsoft, Mozilla and Facebook joined the rust foundation and became Platinum members to promote and develop the rust language around the world. ​

The official website introduces rust: a language that gives everyone the ability to build reliable and efficient software. Rust language has three advantages that deserve your attention:

  1. High performance. Rust is surprisingly fast and memory utilization is extremely high. Because there is no runtime and garbage collection, it can be competent for services with particularly high performance requirements, can run on embedded devices, and can easily integrate with other languages.
  2. Reliability. Rust's rich type system and ownership model ensure memory safety and thread safety, so that you can eliminate all kinds of errors at compile time.
  3. Productivity. Rust has excellent documentation, friendly compiler and clear error prompt information. It also integrates first-class tools - package manager and construction tools, multi editor support for intelligent automatic completion and type verification, and automatic code formatting.

Both Rust and C are hardware direct abstractions
Both rust and C are direct abstractions of hardware and can be regarded as a "portable assembler". Both rust and C can control the memory layout, integer size, stack and heap memory allocation, pointer indirect addressing, etc. of data structures, and can generally be translated into understandable machine code. Compilers rarely insert "magic". Even though rust has higher-level structures than C, such as iterators, trait s, and smart pointers, they are designed to be predictably optimized into simple machine code (also known as "zero cost abstraction"). The memory layout of rust type is very simple. For example, the expandable strings String and Vec are exactly {byte*, capacity, length}. Rust doesn't have any concepts like the move or copy constructor in Cpp, so the passing of objects is guaranteed to be no more complex than passing pointers or memcpy.

In general, Rust has high performance comparable to C, and has efficient development productivity. At the same time, it can efficiently mix programming or call each other with other languages such as C through FFI. This article mainly introduces how to use Rust to develop postgres extension. ​

Deploy development environment
Configure the RUST development environment

reference resources https://www.rust-lang.org/too... IDE recommendations https://code.visualstudio.com/ + https://rls.booyaa.wtf/

Configuring the postgres development environment
ubuntu

sudo apt-get install build-essential libreadline-dev zlib1g-dev flex bison libxml2-dev libxslt-dev libssl-dev libxml2-utils xsltproc

Red hat

sudo yum install -y bison-devel readline-devel zlib-devel openssl-devel wget
sudo yum groupinstall -y 'Development Tools'

cargo pgx subcommand installation
cargo install cargo-pgx

Next, execute the command cargo pgx init.

Discovered Postgres v13.3, v12.7, v11.12, v10.17
  Downloading Postgres v12.7 from https://ftp.postgresql.org/pub/source/v12.7/postgresql-12.7.tar.bz2
  Downloading Postgres v13.3 from https://ftp.postgresql.org/pub/source/v13.3/postgresql-13.3.tar.bz2
  Downloading Postgres v10.17 from https://ftp.postgresql.org/pub/source/v10.17/postgresql-10.17.tar.bz2
  Downloading Postgres v11.12 from https://ftp.postgresql.org/pub/source/v11.12/postgresql-11.12.tar.bz2
     Removing /home/wdy/.pgx/10.17
    Untarring Postgres v10.17 to /home/wdy/.pgx/10.17
  Configuring Postgres v10.17
     Removing /home/wdy/.pgx/11.12
    Untarring Postgres v11.12 to /home/wdy/.pgx/11.12
    Untarring Postgres v12.7 to /home/wdy/.pgx/12.7
  Configuring Postgres v11.12
     Removing /home/wdy/.pgx/13.3
    Untarring Postgres v13.3 to /home/wdy/.pgx/13.3
  Configuring Postgres v12.7
    Compiling Postgres v10.17
  Configuring Postgres v13.3
    Compiling Postgres v11.12
    Compiling Postgres v12.7
    Compiling Postgres v13.3
   Installing Postgres v10.17 to /home/wdy/.pgx/10.17/pgx-install
   Installing Postgres v11.12 to /home/wdy/.pgx/11.12/pgx-install
   Installing Postgres v12.7 to /home/wdy/.pgx/12.7/pgx-install
   Installing Postgres v13.3 to /home/wdy/.pgx/13.3/pgx-install
   Validating /home/wdy/.pgx/10.17/pgx-install/bin/pg_config
 Initializing data directory at /home/wdy/.pgx/data-10
   Validating /home/wdy/.pgx/11.12/pgx-install/bin/pg_config
 Initializing data directory at /home/wdy/.pgx/data-11
   Validating /home/wdy/.pgx/12.7/pgx-install/bin/pg_config
 Initializing data directory at /home/wdy/.pgx/data-12
   Validating /home/wdy/.pgx/13.3/pgx-install/bin/pg_config
 Initializing data directory at /home/wdy/.pgx/data-13

This command will download postgres of versions V10, V11, V12 and V13 and compile it into the directory ~ /. pgx /. This download step is necessary because the subsequent pgx will generate the corresponding Rust bindings for the header file of each version of postgres, and it will also be used in the subsequent pgx test framework. ​

Develop a simple extension
Create an extension project
Use the following command to create a file named my_extension project

cargo pgx new my_extension

After the command is executed, the following directory structure will be generated

├── Cargo.toml
├── my_extension.control
├── sql
│   ├── lib.generated.sql
│   └── load-order.txt
└── src
    └── lib.rs

sql/lib.generated.sql is as follows

CREATE OR REPLACE FUNCTION "hello_my_extension"() RETURNS text STRICT LANGUAGE c AS 'MODULE_PATHNAME', 'hello_my_extension_wrapper';

src/lib.rs is as follows

use pgx::*;

pg_module_magic!();

#[pg_extern]
fn hello_my_extension() -> &'static str {
    "Hello, my_extension"
}

#[cfg(any(test, feature = "pg_test"))]
mod tests {
    use pgx::*;

    #[pg_test]
    fn test_hello_my_extension() {
        assert_eq!("Hello, my_extension", crate::hello_my_extension());
    }

}

#[cfg(test)]
pub mod pg_test {
    pub fn setup(_options: Vec<&str>) {
        // perform one-off initialization when the pg_test framework starts
    }

    pub fn postgresql_conf_options() -> Vec<&'static str> {
        // return any postgresql.conf settings that are required for your tests
        vec![]
    }
}

You can see that pgx has given the simplest extension implementation by default# The function modified by the [pg_extern] macro is the extension function we want to implement, mod tests, pub mod PG_ Test is a test module that pgx has written for us to write relevant test code. ​

pgx has written hello for us by default_ my_ The function of extension is very simple, which is to return the "Hello, my_extension" string

Run extension

cd my_extension 
cargo pgx run pg13  # or pg10 or pg11 or pg12

Use cargo pgx run followed by parameter pg13 or pg10 or pg11 or pg12 to correspond to different Postgres versions. Cargo pgx run will compile the extension into a. so shared library file and copy it to the ~ /. pgx / directory of the corresponding version. Then start the Postgres instance and connect to the database with the same name as the extension through psql. After compilation, the developer will be in the shell interface of psql and can call extension for testing. ​

Here, we execute cargo pgx run pg12, and the output is as follows:

$ cargo pgx run pg12
building extension with features `pg12`
"cargo" "build" "--features" "pg12" "--no-default-features"
    Updating crates.io index
  Downloaded generic-array v0.14.4
  Downloaded humantime v2.1.0
  Downloaded lazycell v1.3.0
  Downloaded shlex v1.0.0
  Downloaded stable_deref_trait v1.2.0
  Downloaded termcolor v1.1.2
  Downloaded typenum v1.13.0
  Downloaded time-macros v0.1.1
  Downloaded which v3.1.1
  Downloaded atomic-traits v0.2.0
  Downloaded seahash v4.1.0
  Downloaded uuid v0.8.2
  Downloaded as-slice v0.1.5
  Downloaded pgx v0.1.21
  Downloaded peeking_take_while v0.1.2
  Downloaded proc-macro-hack v0.5.19
  Downloaded rustc-hash v1.1.0
  Downloaded serde_cbor v0.11.1
  Downloaded time-macros-impl v0.1.2
  Downloaded bindgen v0.58.1
  Downloaded cexpr v0.4.0
  Downloaded env_logger v0.8.4
  Downloaded standback v0.2.17
  Downloaded generic-array v0.12.4
  Downloaded getrandom v0.2.3
  Downloaded clang-sys v1.2.0
  Downloaded glob v0.3.0
  Downloaded const_fn v0.4.8
  Downloaded generic-array v0.13.3
  Downloaded half v1.7.1
  Downloaded hash32 v0.1.1
  Downloaded enum-primitive-derive v0.2.1
  Downloaded heapless v0.6.1
  Downloaded build-deps v0.1.4
  Downloaded libloading v0.7.0
  Downloaded nom v5.1.2
  Downloaded time v0.2.27
  Downloaded pgx-macros v0.1.21
  Downloaded pgx-pg-sys v0.1.21
  Downloaded 39 crates (2.0 MB) in 2.74s
   Compiling version_check v0.9.3
   Compiling libc v0.2.98
   Compiling cfg-if v1.0.0
   Compiling autocfg v1.0.1
   Compiling proc-macro2 v1.0.27
   Compiling unicode-xid v0.2.2
   Compiling syn v1.0.73
   Compiling memchr v2.4.0
   Compiling lazy_static v1.4.0
   Compiling serde_derive v1.0.126
   Compiling tinyvec_macros v0.1.0
   Compiling getrandom v0.1.16
   Compiling matches v0.1.8
   Compiling serde v1.0.126
   Compiling log v0.4.14
   Compiling crossbeam-utils v0.8.5
   Compiling byteorder v1.4.3
   Compiling glob v0.3.0
   Compiling percent-encoding v2.1.0
   Compiling crc32fast v1.2.1
   Compiling ryu v1.0.5
   Compiling adler v1.0.2
   Compiling ppv-lite86 v0.2.10
   Compiling crossbeam-epoch v0.9.5
   Compiling regex-syntax v0.6.25
   Compiling bitflags v1.2.1
   Compiling mime v0.3.16
   Compiling typenum v1.13.0
   Compiling rayon-core v1.9.1
   Compiling scopeguard v1.1.0
   Compiling serde_json v1.0.64
   Compiling unicode-width v0.1.8
   Compiling itoa v0.4.7
   Compiling base64 v0.11.0
   Compiling httpdate v0.3.2
   Compiling xml-rs v0.8.3
   Compiling proc-macro-hack v0.5.19
   Compiling humantime v2.1.0
   Compiling vec_map v0.8.2
   Compiling semver-parser v0.7.0
   Compiling ansi_term v0.11.0
   Compiling unescape v0.1.0
   Compiling bindgen v0.58.1
   Compiling termcolor v1.1.2
   Compiling strsim v0.8.0
   Compiling either v1.6.1
   Compiling shlex v1.0.0
   Compiling lazycell v1.3.0
   Compiling peeking_take_while v0.1.2
   Compiling rustc-hash v1.1.0
   Compiling const_fn v0.4.8
   Compiling stable_deref_trait v1.2.0
   Compiling heapless v0.6.1
   Compiling half v1.7.1
   Compiling cfg-if v0.1.10
   Compiling once_cell v1.8.0
   Compiling seahash v4.1.0
   Compiling libloading v0.7.0
   Compiling tinyvec v1.2.0
   Compiling unicode-bidi v0.3.5
   Compiling unicase v2.6.0
   Compiling nom v5.1.2
   Compiling standback v0.2.17
   Compiling generic-array v0.14.4
   Compiling time v0.2.27
   Compiling memoffset v0.6.4
   Compiling miniz_oxide v0.4.4
   Compiling rayon v1.5.1
   Compiling num-traits v0.2.14
   Compiling form_urlencoded v1.0.1
   Compiling hash32 v0.1.1
   Compiling build-deps v0.1.4
   Compiling clang-sys v1.2.0
   Compiling textwrap v0.11.0
   Compiling semver v0.9.0
   Compiling unicode-normalization v0.1.19
   Compiling rustc_version v0.2.3
   Compiling aho-corasick v0.7.18
   Compiling quote v1.0.9
   Compiling atty v0.2.14
   Compiling dirs-sys v0.3.6
   Compiling socks v0.3.3
   Compiling num_cpus v1.13.0
   Compiling which v3.1.1
   Compiling getrandom v0.2.3
   Compiling crossbeam-channel v0.5.1
   Compiling idna v0.2.3
   Compiling mime_guess v2.0.3
   Compiling atomic-traits v0.2.0
   Compiling generic-array v0.13.3
   Compiling generic-array v0.12.4
   Compiling colored v2.0.0
   Compiling clap v2.33.3
   Compiling rand_core v0.5.1
   Compiling dirs v3.0.2
   Compiling regex v1.5.4
   Compiling uuid v0.8.2
   Compiling flate2 v1.0.20
   Compiling url v2.2.2
   Compiling cexpr v0.4.0
   Compiling as-slice v0.1.5
   Compiling rand_chacha v0.2.2
   Compiling crossbeam-deque v0.8.0
   Compiling env_proxy v0.4.1
   Compiling env_logger v0.8.4
   Compiling rand v0.7.3
   Compiling rttp_client v0.1.0
   Compiling thiserror-impl v1.0.26
   Compiling time-macros-impl v0.1.2
   Compiling enum-primitive-derive v0.2.1
   Compiling time-macros v0.1.1
   Compiling thiserror v1.0.26
   Compiling serde-xml-rs v0.4.1
   Compiling toml v0.5.8
   Compiling serde_cbor v0.11.1
   Compiling pgx-utils v0.1.21
   Compiling pgx-pg-sys v0.1.21
   Compiling pgx-macros v0.1.21
   Compiling pgx v0.1.21
   Compiling my_extension v0.0.0 (/home/wdy/gitlab/valid-my-idea/rust-dig/pgextionsion/my_extension)
    Finished dev [unoptimized + debuginfo] target(s) in 1m 28s

installing extension
      Copying control file to `/home/wdy/.pgx/12.7/pgx-install/share/postgresql/extension/my_extension.control`
      Copying shared library to `/home/wdy/.pgx/12.7/pgx-install/lib/postgresql/my_extension.so`
      Writing extension schema to `/home/wdy/.pgx/12.7/pgx-install/share/postgresql/extension/my_extension--1.0.sql`
     Finished installing my_extension
     Starting Postgres v12 on port 28812
     Creating database my_extension
psql (12.7)
Type "help" for help.

my_extension=#

You can see that at last you enter a file named my_ psql interface of EXTENSION. Next, create EXTENSION and call, and you can see the normal output string

my_extension=# create extension my_extension;
CREATE EXTENSION
my_extension=# select hello_my_extension();
 hello_my_extension  
---------------------
 Hello, my_extension
(1 row)

Array
The Vec type of Rust corresponds to ARRAY []: in postgrs. If an extension function wants to return multiple value s of the same type, it can use Vec to return, such as static in the following example_ Names and static_names_2d, the single row and single column are returned after the psql terminal call. If an extension returns an Iterator, it returns multi row and single column data. See static in the following example_ names_ ITER and static_names_2d_iter. ​

code

#[pg_extern]
fn static_names() -> Vec<Option<&'static str>> {
    vec![Some("King"), Some("Eastern"), None, Some("Sun")]
}

#[pg_extern]
fn static_names_iter() -> impl std::iter::Iterator<Item = Option<&'static str>> {
    vec![Some("Brandy"), Some("Sally"), None, Some("Anchovy")].into_iter()
}

#[pg_extern]
fn static_names_2d() -> Vec<Vec<Option<&'static str>>> {
    vec![
        vec![Some("Brandy"), Some("Sally"), None, Some("Anchovy")],
        vec![Some("Eric"), Some("David")],
        vec![Some("ZomboDB"), Some("PostgreSQL"), Some("Elasticsearch")],
    ]
}

#[pg_extern]
fn static_names_2d_iter() -> impl std::iter::Iterator<Item = Vec<Option<&'static str>>> {
    vec![
        vec![Some("Brandy"), Some("Sally"), None, Some("Anchovy")],
        vec![Some("Eric"), Some("David")],
        vec![Some("ZomboDB"), Some("PostgreSQL"), Some("Elasticsearch")],
    ]
    .into_iter()
}

verification

Execute cargo pgx run pg13

installing extension
      Copying control file to `/home/wdy/.pgx/13.3/pgx-install/share/postgresql/extension/my_extension.control`
      Copying shared library to `/home/wdy/.pgx/13.3/pgx-install/lib/postgresql/my_extension.so`
      Writing extension schema to `/home/wdy/.pgx/13.3/pgx-install/share/postgresql/extension/my_extension--1.0.sql`
     Finished installing my_extension
     Starting Postgres v13 on port 28813
     Re-using existing database my_extension
psql (13.3)
Type "help" for help.

my_extension=# drop extension my_extension;create extension my_extension;
DROP EXTENSION
CREATE EXTENSION
my_extension=# select static_names();
        static_names         
-----------------------------
 {King,Eastern,NULL,Sun}
(1 row)

my_extension=# select static_names_iter();
 static_names_iter 
-------------------
 Brandy
 Sally

 Anchovy
(4 rows)

my_extension=# select static_names_2d();
                                   static_names_2d                                   
-------------------------------------------------------------------------------------
 {"{Brandy,Sally,NULL,Anchovy}","{Eric,David}","{ZomboDB,PostgreSQL,Elasticsearch}"}
(1 row)

my_extension=# select static_names_2d_iter();
        static_names_2d_iter        
------------------------------------
 {Brandy,Sally,NULL,Anchovy}
 {Eric,David}
 {ZomboDB,PostgreSQL,Elasticsearch}
(3 rows)

Complete code
Complete code

AGGREGATE extension
Postgres can define a new aggregation function by creating aggregate. A simple aggregate function contains one or two ordinary functions:

  1. State transition function sfunc
  2. Optional final calculation function ffunc
sfunc( internal-state, next-data-values ) ---> next-internal-state
ffunc( internal-state ) ---> aggregate-value
  1. You may also need to define the data type stype for the internal state of the aggregate function
  2. The initial value of the internal function state, initcond, is a value of type text.

Then we only need to refer to the above example to develop the corresponding sfunc, ffunc and stype, and then use CREATE AGGREGATE to create the aggregation function. ​

Create extension
Execute the following command
cargo pgx new my_agg

Define internal state type stype

The type IntegerAvgState is used to store intermediate state data during aggregation function calculation

#[derive(Serialize, Deserialize, PostgresType)]
pub struct IntegerAvgState {
    sum: i32,
    n: i32,
}
impl Default for IntegerAvgState {
    fn default() -> Self {
        Self { sum: 0, n: 0 }
    }
}
impl IntegerAvgState {
    fn acc(&self, v: i32) -> Self {
        Self {
            sum: self.sum + v,
            n: self.n + 1,
        }
    }
    fn finalize(&self) -> i32 {
        self.sum / self.n
    }
}

State transition function sfunc

#[pg_extern]
fn integer_avg_state_func(
    internal_state: IntegerAvgState,
    next_data_value: i32,
) -> IntegerAvgState {
    internal_state.acc(next_data_value)
}

Final calculation function ffunc

#[pg_extern]
fn integer_avg_final_func(internal_state: IntegerAvgState) -> i32 {
    internal_state.finalize()
}

Create aggregate function

extension_sql!(
    r#"
    CREATE AGGREGATE MYAVG (integer)
    (
        sfunc = integer_avg_state_func,
        stype = IntegerAvgState,
        finalfunc = integer_avg_final_func,
        initcond = '{"sum": 0, "n": 0}'
    );
    "#
);

Complete code

Validate aggregate function
Execute cargo pgx run pg12 to enter the psql command line interface of postgres, and then perform the following operations

my_agg=# DROP EXTENSION my_agg; CREATE EXTENSION my_agg;
DROP EXTENSION
CREATE EXTENSION
my_agg=#  create table t (c integer);
CREATE TABLE
my_agg=# insert into t (c) values (1), (2), (3);
INSERT 0 3
my_agg=# select MYAVG(c) from t;
 myavg 
-------
     2
(1 row)

my_agg=# drop table t;
DROP TABLE
my_agg=#  create table t (c integer,b text);
CREATE TABLE
my_agg=# insert into t (c,b) values (1,'king'), (2,'eastern'), (3,'sun');
INSERT 0 3
my_agg=# select MYAVG(c) from t;
 myavg 
-------
     2
(1 row)

Complete code
Complete code

TOPN AGGREGATE extension
Now we implement a slightly more complex aggregation function to find the maximum 10 values of a column of value in the table. The classical method is to use the minimum heap. BinaryHeap in the Rust standard library is the maximum heap by default. You can get a minimum heap by encapsulating it with Reverse. ​

Create extension

cargo pgx new my_topn

Define internal state type stype

#[derive(Serialize, Deserialize, PostgresType)]
pub struct TopState {
    min_heap:  BinaryHeap<Reverse<i32>>,
    n:usize,
}
impl Default for TopState {
    fn default() -> Self {
        Self { min_heap:BinaryHeap::new(),n:10 }
    }
}

Implement acc method for TopState for state transition and update the minimum heap

impl TopState {
    fn acc(& mut self, v: i32)  {
        if self.min_heap.len()<self.n{
            self.min_heap.push(Reverse(v));
            return 
        }

        // Gets the minimum value on the current heap
        let top = self.min_heap.peek().unwrap().0;
        // If it is smaller than the minimum value, it will not be the maximum 10 values required, and it will be discarded directly
        if v<=top{
            return 
        }

        // Insert into the minimum heap and remove the minimum value from the heap
        self.min_heap.push(Reverse(v));
        self.min_heap.pop().unwrap();
        return 

    }

}

State transition function sfunc

#[pg_extern]
fn integer_topn_state_func(
   mut internal_state: TopState,
    next_data_value: i32,
) -> TopState {
    internal_state.acc(next_data_value);
    internal_state
}

Final calculation function ffunc

#[pg_extern]
fn integer_topn_final_func(internal_state: TopState) -> Vec<i32> {
    internal_state.min_heap.into_sorted_vec().iter().map(|x|x.0).collect()
}

Create aggregate function MYMAXN

extension_sql!(
    r#"
    CREATE AGGREGATE MYMAXN (integer)
    (
        sfunc = integer_topn_state_func,
        stype = TopState,
        finalfunc = integer_topn_final_func,
        initcond = '{"min_heap":[],"n":10}'
    );
    "#
);

Validate aggregate function MYMAXN

Execute cargo pgx run pg13 to enter the psql command line interface of postgres, and execute the following sql command to create test data

create table people
(
    id        integer,
    name      varchar(32),
    age       integer,
    grade     numeric(4, 2),
    birthday  date,
    logintime timestamp
);



insert into people
select generate_series(1,20) as id,
md5(random()::text) as name,
(random()*100)::integer as age,
(random()*99)::numeric(4,2) as grade,
now() - ((random()*1000)::integer||' day')::interval as birthday,
clock_timestamp() as logintime;

The test data are as follows

my_topn=# select age from people ;
 age 
-----
  28
  62
   3
  20
  46
  23
  74
  19
  46
  26
  70
  90
  22
  45
  30
  46
  43
  70
  78
  96
(20 rows)

Execute aggregate function MYMAXN

my_topn=# select MYMAXN(age) from people;
             mymaxn              
---------------------------------
 {96,90,78,74,70,70,62,46,46,46}
(1 row)

Complete code

be careful
If initcond is not specified, the following error prompt will be reported when creating the extension

ERROR:  must not omit initial value when transition function is strict and transition type is not compatible with input type

Upgraded TOPN

Return the maximum n values and the minimum n values of the specified column in the table at the same time. Add a maximum heap to the original data structure. ​

Define internal state type stype

#[derive(Serialize, Deserialize, PostgresType)]
pub struct TopState {
    min_heap:  BinaryHeap<Reverse<i32>>,
    max_heap:  BinaryHeap<i32>,
    n:usize,
}
impl Default for TopState {
    fn default() -> Self {
        Self { min_heap:BinaryHeap::new(),max_heap:BinaryHeap::new(),n:10 }
    }
}

Implement acc method for TopState for state transition and update the minimum heap

impl TopState {
    fn acc(& mut self, v: i32)  {
        if self.min_heap.len()<self.n{
            self.min_heap.push(Reverse(v));
            return 
        }

        // Gets the minimum value on the current heap
        let top = self.min_heap.peek().unwrap().0;
        // If it is smaller than the minimum value, it will not be the maximum 10 values required, and it will be discarded directly
        if v<=top{
            return 
        }

        // Insert into the minimum heap and remove the minimum value from the heap
        self.min_heap.push(Reverse(v));
        self.min_heap.pop().unwrap();
        return 

    }
    fn acc_max(& mut self, v: i32)  {
        if self.max_heap.len()<self.n{
            self.max_heap.push(v);
            return 
        }

        // Fetches the maximum value on the current heap
        let top = self.max_heap.peek().unwrap();
        // If it is larger than the maximum value, it will not be the minimum 10 values required, and it will be discarded directly
        if v>=*top{
            return 
        }

        // Insert into the maximum heap and remove the maximum value from the heap
        self.max_heap.push(v);
        self.max_heap.pop().unwrap();
        return 

    }

}

State transition function sfunc

#[pg_extern]
fn integer_topn_state_func(
   mut internal_state: TopState,
    next_data_value: i32,
) -> TopState {
    internal_state.acc(next_data_value);
    internal_state.acc_max(next_data_value);
    internal_state
}

Final calculation function ffunc

#[pg_extern]
fn integer_topn_final_func(internal_state: TopState) -> Vec<Vec<i32>> {
    vec![
    internal_state.min_heap.into_sorted_vec().iter().map(|x|x.0).collect(),
    internal_state.max_heap.into_sorted_vec(),
    ]
}

Create aggregate function MYMAXN

extension_sql!(
    r#"
    CREATE AGGREGATE MYMAXN (integer)
    (
        sfunc = integer_topn_state_func,
        stype = TopState,
        finalfunc = integer_topn_final_func,
        initcond = '{"min_heap":[],"max_heap":[],"n":10}'
    );
    "#
);

Validate aggregate function MYMAXN

my_topn=# select MYMAXN(age) from people;
                                mymaxn                                
----------------------------------------------------------------------
 {"{96,90,78,74,70,70,62,46,46,46}","{3,19,20,22,23,26,28,30,43,45}"}
(1 row)

Complete code

reference

  1. https://mp.weixin.qq.com/s/9r...
  2. Postgres Compile and Install from source code https://wiki.postgresql.org/w...
  3. Introduction to cargo pgx command https://github.com/zombodb/pg...
  4. https://blog.timescale.com/bl...
  5. https://www.postgresql.org/do...

Topics: Database Rust