pg_bigmをRustで実装する（Part1）
この記事はPostgreSQL Advent Calendar 2024の19日目の記事です。
昨日は澤田さんの「PostgreSQL 17で新しく実装されたradix treeを使ってインメモリのキーバリューストア作ってみた」でした。拡張機能を実装してさらにPostgreSQLのバグを見つけているのがさすがだなと思いました。
はじめに
2024年10月にギリシャで開催されたPGConf.EU 2024に参加したのですが、そのときに「Fearless Extension Development With Rust and PGRX」という発表がありました。これは、PostgreSQLのエクステンションをRustで実装するためのフレームワークであるpgrxに関する発表でした。
すごく面白い発表だったので、自分でもRustでPostgreSQLのエクステンションを実装してみます。今回は、よく使うエクステンションであるpg_bigmをRustで実装します。が、実装が終わらなかったので、この記事はPart1として、Part2以降はいずれ書きます。
ソースコードはGitHubで公開しています。実装のモチベーションが上がるのでStar付けていただけると嬉しいです。
pgrx
pgrxの特徴
PostgreSQLのエクステンションは基本的にはC言語で実装しますが、それをRustで実装するためのフレームワークがpgrxになります。pgrxの特徴は以下です。
- 完全管理された開発環境：cargo-pgrxを使用して、PostgreSQL拡張機能の開発を迅速に行うことができる。
- 複数のPostgreSQLバージョンのサポート：PostgreSQL 12から17までのバージョンを同一のコードベースでサポートし、バージョン特有のAPIをRustの機能ゲーティングで使用できる。
- 自動スキーマ生成：Rustで拡張機能を完全に実装し、多くのRust型をPostgreSQLに自動的にマッピングする。
- 安全性の優先：Rustのpanic!をPostgreSQLのERRORに変換し、トランザクションを中断させることなく安全にメモリ管理を行う。
-
UDFサポート：関数に
#[pg_extern]を付与することでPostgreSQLに公開し、SETOFやTABLEを返すことができる。 - 簡単なカスタムタイプ：Rustの構造体や列挙型をPostgreSQLの型として使用でき、カスタム表現を提供する。
-
サーバープログラミングインターフェース（SPI）：SPIへの安全なアクセスを提供し、所有権のある
DatumをSPIコンテキストから透過的に返す。 - 高度な機能：PostgreSQLのメモリコンテキストシステムへの安全なアクセスや、PostgreSQLの内部への直接的なunsafeアクセスを可能にします。
pgrxを使っているエクステンション
また、pgrxを使って開発されているエクステンションとして有名なものは以下があります。
つらつらと書きましたが、個人的に一番いいところだと思ったのは、環境構築の容易さです。
pgrxを使ったエクステンション開発環境構築
- 事前にSystem Requirementsを参考に、RustやPostgreSQLコンパイルに必要なパッケージをインストール
- cargo-pgrxをインストール
- 筆者の環境では1分程度かかりました。
$ cargo install --locked cargo-pgrx
- pgrxホームディレクトリを初期化
- PostgreSQL 12-17をすべてコンパイルするので時間がかかります。筆者の環境では6分程度かかりました。
- デフォルトではホームディレクトリ配下に
.pgrxディレクトリが作成され、その中にPostgreSQLのソースコードやデータディレクトリが作成されます。
$ cargo pgrx init
$ ls -l ~/.pgrx
total 76
drwxrwxrwx 7 shinya shinya 4096 Dec 13 15:24 12.22
drwxrwxr-x 2 shinya shinya 4096 Dec 13 15:21 12.22_unpack
drwxrwxrwx 7 shinya shinya 4096 Dec 13 15:24 13.18
drwxrwxr-x 2 shinya shinya 4096 Dec 13 15:21 13.18_unpack
drwxrwxrwx 7 shinya shinya 4096 Dec 13 15:24 14.15
drwxrwxr-x 2 shinya shinya 4096 Dec 13 15:20 14.15_unpack
drwxrwxrwx 7 shinya shinya 4096 Dec 13 15:24 15.10
drwxrwxr-x 2 shinya shinya 4096 Dec 13 15:20 15.10_unpack
drwxrwxrwx 7 shinya shinya 4096 Dec 13 15:24 16.6
drwxrwxr-x 2 shinya shinya 4096 Dec 13 15:21 16.6_unpack
drwxrwxr-x 7 shinya shinya 4096 Dec 13 15:25 17.2
drwxrwxr-x 2 shinya shinya 4096 Dec 13 15:21 17.2_unpack
-rw-rw-r-- 1 shinya shinya 374 Dec 13 15:25 config.toml
drwx------ 19 shinya shinya 4096 Dec 13 15:25 data-12
drwx------ 19 shinya shinya 4096 Dec 13 15:25 data-13
drwx------ 19 shinya shinya 4096 Dec 13 15:25 data-14
drwx------ 19 shinya shinya 4096 Dec 13 15:25 data-15
drwx------ 19 shinya shinya 4096 Dec 13 15:25 data-16
drwx------ 19 shinya shinya 4096 Dec 13 15:25 data-17
- pgrxプロジェクトを作成
- 今回はpg_bigmをRustで実装するので、pg_bigmrという名前にします。
$ cargo pgrx new pg_bigmr
- エクステンションをコンパイルして実行
- なんと
cargo pgrx runするだけでエクステンションのコンパイルができます。そして、そのPostgreSQLに接続されます。まだ1行もプログラムを書いていないですが、「Hello, pg_bigmr」という文字列を出力するユーザ定義関数hello_pg_bigmr()が定義されています。簡単ですね。
- なんと
$ cd pg_bigmr
$ cargo pgrx run
~snip~
pg_bigmr=# CREATE EXTENSION pg_bigmr;
CREATE EXTENSION
pg_bigmr=# \dx+ pg_bigmr
Objects in extension "pg_bigmr"
Object description
---------------------------
function hello_pg_bigmr()
(1 row)
pg_bigmr=# SELECT hello_pg_bigmr();
hello_pg_bigmr
-----------------
Hello, pg_bigmr
(1 row)
pg_bigmrの実装
それでは、ここからpg_bigmをRustで実装していきますが、その前にpg_bigmを実装するには何を実装すればいいかをおさらいしておきます。
pg_bigmとは？
PostgreSQL上で全文検索機能を提供するモジュールです。詳しくはpg_bigmの作者である澤田さんのスライドが参考になります。
ソースコードはこちらです。
pg_bigmの構成要素
現状では、「Part1で実装するもの」に◯をしたオブジェクトの実装が完了しています。◎のものを本記事で取り上げます。その他はソースコードをご参照ください。
| 種類 | オブジェクト名 | Part1で実装するもの |
|---|---|---|
| GUCパラメータ | pg_bigm.enable_recheck |
◎ |
| GUCパラメータ | pg_bigm.gin_key_limit |
◯ |
| GUCパラメータ | pg_bigm.similarity_limit |
◯ |
| GUCパラメータ | pg_bigm.last_update |
◯ |
| 演算子 | =% |
◎ |
| 演算子クラス | gin_bigm_ops |
◎ |
| 演算子族 | gin_bigm_ops |
◎ |
| 関数（ユーザ用） | likequery() |
◎ |
| 関数（ユーザ用） | show_bigm() |
◯ |
| 関数（ユーザ用） | bigm_similarity() |
◯ |
| 関数（ユーザ用） | pg_gin_pending_stats() |
|
| 関数（GINインデックス用） | bigmtextcmp() |
|
| 関数（GINインデックス用） | gin_extract_value_bigm() |
|
| 関数（GINインデックス用） | gin_extract_query_bigm() |
|
| 関数（GINインデックス用） | gin_bigm_consistent() |
|
| 関数（GINインデックス用） | gin_bigm_compare_partial() |
|
| 関数（GINインデックス用） | gin_bigm_triconsistent() |
GUCパラメータ（pg_bigm.enable_recheck）
関数名こそ違いますが、同様の感じでGUCパラメータを定義できます。
pg_bigm
- 関数名は
DefineCustomBoolVariable() - デフォルト値は、
DefineCustomBoolVariable()の引数で設定 -
long_desc（詳細なパラメータについての説明文）にNULLを設定可能 -
check_hook、assign_hook、how_hook（後半3つの引数）を設定可能
pg_bigmr
- 関数名は
define_bool_guc() - デフォルト値は、
define_bool_guc()の引数ではなく、static変数の初期値として設定 -
long_description（詳細なパラメータについての説明文）の型が&strなのでNoneが設定できない（RustではNoneとなり得る値はOptionで包む必要がある）ので、今回は空文字列（""）を設定 -
check_hook、assign_hook、how_hookが設定できなさそう。詳細は未調査
演算子（=%）
Cでエクステンションを書くときには、SQLでのCREATE OPERATOR文とそのための関数を実装する必要があります。一方で、Rustでエクステンションを書くときには、Rustの関数にアトリビュートを設定しておくと、pgrxが自動でSQL文を生成してくれるため、SQLを実装する必要はありません。
pg_bigm
- SQLとCの関数の実装が必要
pg_bigmr
- Rustの関数にアトリビュートを設定すれば、pgrxがSQLを自動生成
-
#[pg_operator(parallel_safe, stable, strict)]：関数のパラメータ -
#[opname(=%)]：演算子名 -
#[restrict(contsel)]：演算子のパラメータ -
#[join(contjoinsel)]：演算子のパラメータ
演算子クラス、演算子族（gin_bigm_ops）
演算子クラス、演算子族の定義には、C、RustともにSQLとそのための関数の実装が必要です。演算子と同様にSQL文の自動生成に対応しているとよかったのですが、現状では、B-treeインデックスやHashインデックスにしか対応しておらず、GINインデックスには対応していないようです。
pg_bigm
- SQLの実装が必要
- 演算子クラス、演算子族には演算子、関数の実装も必要ですが、ここでは省略
pg_bigmr
-
CREATE EXTENSION時に実行されるSQLをextension_sql_file!()マクロで指定-
extension_sql!()マクロでソースコード内にSQLを直接書くこともできる -
ALTER EXTENSION時に実行されるSQLは、sqlディレクトリにpg_bigmr--0.1.0--0.1.1.sqlの形式で格納しておけばよい
-
関数（ユーザ用）（likequery）
Cでエクステンションを書くときには、SQLでのCREATE FUNCTION文とその実装が必要です。一方で、Rustでエクステンションを書くときには、Rustの関数にアトリビュートを設定しておくと、pgrxが自動でSQL文を生成してくれるため、SQLを実装する必要はありません。
pg_bigm
- SQLとその実装が必要
pg_bigmr
- Rustの関数にアトリビュートを設定すれば、pgrxがSQLを自動生成
-
#[pg_extern(immutable, parallel_safe, strict)]：関数のパラメータ
難しかった点
- そもそもRustを書いたことがなかったので、Rustの勉強に苦労した
- ドキュメントがあまりないので、他のエクステンションのソースコードを読みつつ実装進めたので、時間がかかった
- C（PostgreSQL）とRustの様々なフォーマットの違いを理解するのに苦労した
- Cの文字列はヌル終端(
\0)する必要があるが、Rustではそうではない - PostgreSQLのvarlenaとRustのStringのメモリレイアウトが違う
- Cの文字列はヌル終端(
- Rust側で確保したメモリをRust側で解放せずにCに渡したり、Cで確保したメモリをRust側から参照したりなど、CとRust間のやり取りを考慮するのが大変だった
今後の予定
- 残オブジェクトの実装
- テストの実装
- 性能評価
- ほぼ確実にCのほうが早いと思いますが、どのくらいの差が出るのかを確認したい
まとめ
RustでPostgreSQLのエクステンションを実装するためのフレームワークであるpgrxを使って、pg_bigmをRustで実装しました。
SQLの自動生成であったり、Rustの開発ツール（Cargo）を使って開発できる点などが良かったですが、「難しかった点」にも書いたように、実装するにあたってドキュメントがあまりないので、思ったよりも実装に時間がかかってしまいました。
明日は篠田さんの「PostgreSQLとOracle Databaseでシーケンスの動作を比較」です。私はOracleを使ったことがないので、どのような違いがあるか楽しみですね！
Discussion