RustでPostgreSQLのユーザー定義関数を書く

PostgreSQLのユーザー定義関数をRustで実装する話です。

こんにちは、エンジニアの松本です。主な業務としてインメモリデータベースをRustで実装しています。

フォルシアではPostgreSQLを使っており、C言語で拡張も書いていますが、Rustを使って書けるようになると環境構築やテストがしやすくなって嬉しいです。本記事ではRustで関数を実装するとPostgreSQLから使えるようにラップしてくれるzombodb/pgxというクレートを紹介します。 C言語実装との比較実験を行い、遜色ない速度で実行できることを確認しました。

環境構築

環境はUbuntu 20.04.1 LTS (Focal Fossa)で行います。

PostgreSQL13.1を公式の手順でインストールしました。加えてsudo ln -s /usr/local/pgsql-13.1 /usr/local/pgsql とシンボリックリンクを張り、export PATH=/usr/local/pgsql/bin:$PATH としてパスを通している状態です。

何をやるか

SQLでは扱いにくい処理を行うときにユーザー定義関数を書くことが多いです。

ループを含むような処理の例として、コラッツ予想で知られている「整数nについて偶数ならば n = n/2 、奇数ならば n = 3*n+1 とする」という手順を 「n == 1となるまで繰り返すときの回数」を返すような関数を作ります。

C言語で書く場合

まずはC言語での実装を示します。詳細はドキュメントを参照して下さい。

# collatz.c
#include "fmgr.h"
#include "postgres.h"
// `int32`, `int64` は `postgres.h` 内で定義されている。

PG_MODULE_MAGIC;

PG_FUNCTION_INFO_V1(collatz_c);
Datum collatz_c(PG_FUNCTION_ARGS);

Datum collatz_c(PG_FUNCTION_ARGS) {
  int32 arg = PG_GETARG_INT32(0);  // 第1引数をint32として取得する
  int64 n = arg;
  int32 count = 0;
  while (n > 1) {
    if (n % 2 == 0) {
      n /= 2;
    } else {
      n = 3 * n + 1;
    }
    count += 1;
  }
  PG_RETURN_INT32(count);  // countをint32として返却する
}

コンパイルを行い、

$ gcc -shared -O2 -Wall -fpic -I/usr/local/pgsql/include/server collatz.c  -o collatz_c.so
$ sudo mv collatz_c.so /usr/local/pgsql/lib/

# CREATE or REPLACE FUNCTION collatz_c(int4) RETURNS int4 AS 'collatz_c.so', 'collatz_c' LANGUAGE C IMMUTABLE STRICT;
CREATE FUNCTION

# select collatz_c(12);
 collatz_c 
-----------
         9

12, 6, 3, 10, 5, 16, 8, 4, 2, 1 と遷移するので出力 9 が正しいことが確認できます。

pgxを使ってRustで書く場合

本題であるzombodb/pgxを紹介します。PostgreSQL 10~13に対応しています。

cargo install cargo-pgx でサブコマンドをインストールします。 cargo pgx init を実行すると、pgxの検証用にPostgreSQL 10~13の各バージョンがインストールされます。ご飯が食べられるくらいには時間がかかります。

cargo pgx collatz として、ボイラーテンプレートからプロジェクトを作成します。 src/lib.rs に処理を実装します。

# lib.rs
use pgx::*;
pg_module_magic!();

#[pg_extern(immutable)]
fn collatz_strict(arg: i32) -> i32 {
    if arg <= 0 {
        panic!("The function 'collatz' expected the arg is a positive integer.")
    }
    let mut count = 0;
    let mut n = arg as i64;
    while n > 1 {
        n = if n % 2 == 0 { n / 2 } else { 3 * n + 1 };
        debug_assert!(n >= 1);
        count += 1;
    }
    count
}

#[pg_extern(immutable)]
fn collatz(arg: Option) -> i32 {
    match arg {
        Some(arg) => {
            if arg <= 0 {
                panic!("The function 'collatz' got a not positive integer, expected the arg is a positive integer.")
            }
            let mut count = 0;
            let mut n = arg as i64;
            while n > 1 {
                n = if n % 2 == 0 { n / 2 } else { 3 * n + 1 };
                debug_assert!(n >= 1);
                count += 1;
            }
            count
        }
        None => {
            panic!("The function 'collatz' got a null, expected the arg is a positive integer.")
        }
    }
}

#[test]
fn test_collatz() {
    assert_eq!(0, collatz_strict(1));
    assert_eq!(1, collatz_strict(2)); // 2,1
    assert_eq!(7, collatz_strict(3)); // 3,10,5,16,8,4,2,1
    assert_eq!(2, collatz_strict(4)); // 4,2,1
    assert_eq!(5, collatz_strict(5)); // 5,16,8,4,2,1
    assert_eq!(8, collatz_strict(6)); // 6,3,10,5,16,8,4,2,1
    assert_eq!(16, collatz_strict(7)); // 7,22,11,34,17,52,26,13,40,20,10,5,...,1
}
#[test]
#[should_panic]
fn test_collatz_panic() {
    collatz_strict(0);
}

cargo test で動作確認をすることができます。同じファイルに手軽にテストを書き、標準のパッケージマネージャから実行できる点はRustの長所の一つだと感じます。

cargo pgx package でリリースビルドを行うと、target/release 以下に必要なファイル群が作成されます。

$ tree target/release/collatz-pg13/usr/local/pgsql-13.1/
target/release/collatz-pg13/usr/local/pgsql-13.1/
├── lib
│   └── collatz.so
└── share
    └── extension
        ├── collatz--1.0.sql
        └── collatz.control

collatz--1.0.sql を確認すると、 collatz_strict には STRICT をつけて宣言していることが確認できます。引数に Option 型が含まれない場合は自動で STRICT をつけた宣言が作成されるようになっています。

-- collatz--1.0.sql
CREATE OR REPLACE FUNCTION "collatz_strict"("arg" integer) RETURNS integer STRICT IMMUTABLE LANGUAGE c AS 'MODULE_PATHNAME', 'collatz_strict_wrapper';
CREATE OR REPLACE FUNCTION "collatz"("arg" integer) RETURNS integer IMMUTABLE LANGUAGE c AS 'MODULE_PATHNAME', 'collatz_wrapper';

必要なファイルを移動し、extensionとして登録します。

$ sudo mv target/release/collatz-pg13/usr/local/pgsql-13.1/lib/collatz.so /usr/local/pgsql/lib/
$ sudo mv target/release/collatz-pg13/usr/local/pgsql-13.1/share/extension/collatz* /usr/local/pgsql/share/extension/

-- create extension collatz;
CREATE EXTENSION

select collatz(12);
 collatz 
---------
       9

速度比較

作成したそれぞれの関数について100万回実行時の速度を検証します。

-- テスト用テーブルを作成
# create table numbers as (select generate_series(1,1000000) as num);
SELECT 1000000
Time: 1210.538 ms

-- 結果が等しいことを確認
# select * from (select num, collatz(num) as rust, collatz_c(num) as c from numbers)s where rust!=c ;
num | rust | c 
-----+------+---
(0 rows)

-- 速度検証用のコマンド
# select sum(collatz_strict(num)) from numbers ;
    sum    
-----------
 131434424
(1 row)

# select sum(collatz_c(num)) from numbers ;
    sum    
-----------
 131434424
(1 row)

各関数について3回実行したところ下記の結果となりました。

関数	1回目[ms]	2回目[ms]	3回目[ms]
C(collatz_c)	464.390	475.266	465.974
Rust(collatz_strict)	455.088	449.032	461.443
Rust(collatz)	446.275	442.574	451.141

あくまで私の環境での実測値になりますが、Rust実装の方がC実装よりも高速に処理されることが確認できました。環境構築やテストの利便性を考えればRustに移行したほうがよいと考えられます。

まとめ

本記事ではPostgreSQLのユーザー定義関数をpgxを使ってRustで実装しました。C言語実装の関数と比べて遅くないことを確認しました。

本記事では紹介しきれませんでしたが、pgxには一通り必要な機能が揃っているように感じました。新しい関数は当然Rustで書くよねという時代が来るかもしれません。

フォルシアではPostgreSQLのパフォーマンス改善に強いエンジニアを募集しています。