RustでPostgreSQLの拡張を作ってみる

Rust 初心者のsonesuke( https://twitter.com/sonesuke )です。
これは、Rust Advent Calendar 2021 の 16日目です。遡って書いています。

仕事でPostgreSQL拡張機能を作るシチュエーションが出てきたのですが、「C言語で拡張つくるのもつまらないな」と思っていたら、Rustでも作れるみたいだったので、作ってみました。

TL;DR

1. PostgreSQLに日本語形態素解析をする拡張を作ってみた
2. cargo-pgxを使ってRustで簡単にPostgreSQLを作ることができる
3. 成果物は以下

要件定義

真面目に、要件定義してみます。

1. 日本語の形態素解析がしたい
2. 解析対象はPostgreSQLのTEXTカラム
3. インストールが簡単 → PostgreSQLにコピーしたら使える (ビルドし直さない)
4. パフォーマンスがよいこと
5. 筆者の知的好奇心を満たす

技術選定してみる

いろいろ探します・・・・・

うーん、筆者がPythonをかけるので、簡単そうなんだけどFDWだし・・・Pythonだし・・・パフォーマンス悪そうだし・・・
まだまだ、探す・・・・

真面目に、作るしかないのか、パフォーマンスといえば、C言語でネイティブっしょ。

Datum・・・・PG_FUNCTION_ARGS・・・・ふむふむ・・・・ネイティブで、作るのは思ったよりめんどくさいな。何より、筆者の知的好奇心を満たさない。却下。
めげずに、探す・・・

あった・・・・cargo-pgx・・・・どうやら、Rustで簡単に作れるらしい。

Rustはやってみたいと思いつつも、ずっと放置していました。これなら知的好奇心も満たせます。

全ての要件を満たせました！

cargo-pgxってなに？

PostgreSQLを Rustで簡単に作るためのcrateです。
これを使うと、インストールして数コマンドを打つだけで、Rustで拡張を作れるとのこと・・・英語がわかる人は、下記の動画を見れば、どれだけ簡単かわかるかと・・・・ 筆者は雰囲気しか見てない

これを使っていくことにしましょう。

用意するもの

1. Rustが使える環境
2. cargo-gpx
3. 楽しむこころ

セットアップ

Rustが使える環境を用意したら、cargo-pgxをセットアップしていきます。

$ cargo install cargo-pgx

チュートリアルに従って、初期化します。

$ cargo pgx init

サポートしている数のPostgreSQLのビルド祭が始まります。とても時間かかります。筆者の環境では、10から14までの5バージョンのビルドが入りました。

次に、拡張のプロジェクト作成します。

$ cargo pgx new ja_tokenizer

これで、テンプレが作成されます。できたら、試してみましょう。

$ cd ja_tokenizer
$ cargo pgx run pg14

しばらくすると、psqlが立ち上がるので、できたものを試すことができます。

psql (14.1)
Type "help" for help.

ja_tokenizer=# CREATE EXTENSION ja_tokenizer;
CREATE EXTENSION
ja_tokenizer=# SELECT hello_ja_tokenizer();
 hello_ja_tokenizer
---------------------
 Hello, ja_tokenizer
(1 row)

それっぽいですね。C言語で拡張を作っていくよりは圧倒的に楽そうです。あとは、実装あるのみ！

Rustで日本語の形態素解析

形態要素解析としては、linderaを使用します。

選定理由は、Rustで完結するから(mecab-rsに挫折したから)

まずはCargo.tomlに依存関係を記載して、サンプルみながら、さらっと実装しましょう。

#[pg_extern]
fn ja_tokenize(input: &'static str) -> Vec<String> {
    let mut tokenizer = Tokenizer::new().unwrap();

    let tokens = tokenizer.tokenize(input).unwrap();

    let mut ret = Vec::<String>::new();
    for token in tokens {
        ret.push(String::from(token.text));
    }
    ret
}

早速、psqlで試してみます。

ja_tokenizer=# DROP EXTENSION ja_tokenizer;
DROP EXTENSION
ja_tokenizer=# CREATE EXTENSION ja_tokenizer;
CREATE EXTENSION
ja_tokenizer=# SELECT ja_tokenize('新しい日本の夜明け');
       ja_tokenize
-------------------------
 {新しい,日本,の,夜明け}
(1 row)

よし、それっぽい。
DROP EXTENSION をしないとリロードされなくて、ハマりました。
うーん。ここでビルドが一気に重たくなったぞ。。。。これについては、宿題にします。

複数のレコードを返したい

それっぽくできたのですが、1つのレコードを、全結果を返すと、その後SQLも書きづらいです。全部をArrayで返すのではなくて、複数のレコードに分けて、返すことはできないでしょうか？

まじめに、cargo-pgxのドキュメントを読んでみます。

Return impl std::iter::Iterator<Item = T> where T: IntoDatum for automatic set-returning-functions (both RETURNS SETOF and RETURNS TABLE (...) variants

どうやら、Iterartorを返せば、複数レコードになる模様。exampleをあさってみると、それっぽいのがあるので、以下を参考にやってみます。

#[pg_extern]
fn split_set(
    input: &'static str,
    pattern: &'static str,
) -> impl std::iter::Iterator<Item = &'static str> {
    input.split_terminator(pattern).into_iter()
}

exampleが充実していて、よいCrateですね。助かります。
先のコードの戻り値を工夫します。

#[pg_extern]
fn jat_tokenize(input: &str) -> impl std::iter::Iterator<Item = String> {
    let mut tokenizer = Tokenizer::new(Mode::Normal, "");
    let tokens = tokenizer.tokenize(input);
    let mut ret = Vec::<String>::new();
    for token in tokens {
        ret.push(String::from(token.text));
    }
    ret.into_iter()
}

あらためて、psqlで試してみます。

psql (14.1)
Type "help" for help.

ja_tokenizer=# DROP EXTENSION ja_tokenizer;
DROP EXTENSION
ja_tokenizer=# CREATE EXTENSION ja_tokenizer;
CREATE EXTENSION
ja_tokenizer=# SELECT ja_tokenize('新しい日本の夜明け');
 ja_tokenize
-------------
 新しい
 日本
 の
 夜明け
(4 rows)

チューニング

今のままでは、linderaの初期化がレコードを処理するごとに呼ばれてしまいます。明らかに効率が悪そうです。

初期化を一回にしたいのですが・・・・C言語でいうと、これをstaticにするのですが、Rustではどうなるか？
この辺を読むと、今時は、once_cellというものを使うらしい・・・cargo-pgxの中でもonce_cellがたくさん使われているので、きっとあってると願いつつ。。。。

コンパイラーに怒られながら、それっぽく実装・・・・(実はここが一番苦労しました。)

static TOKENIZER: OnceCell<Mutex<Tokenizer>> = OnceCell::new();

#[pg_extern]
fn ja_tokenize(input: &str) -> impl std::iter::Iterator<Item=String> {
    let t = TOKENIZER.get_or_init(|| Mutex::new(Tokenizer::new().unwrap()));
    let result = t.lock().unwrap().tokenize(input);
    let mut ret = Vec::<String>::new();
    match result {
        Err(why) => panic!("{:?}", why),
        Ok(tokens) => {
            for token in tokens {
                ret.push(String::from(token.text));
            }
        }
    }
    ret.into_iter()
}

正直、まだ、理解ができてないないながら、とりあえずクリア。

まとめ

便利な時代になりましたね。
Rust初心者が、勉強しながら、PostgreSQL拡張を作れるようになるまで半日でした。

それでは、良いお年を！

積み残し

以下、積み残しなので、暇を見つけてやってみようかと思っています。

• ユーザ辞書を使いたい (解説記事はこちら)
• エラー処理を真面目にやってない
• ドキュメンテーションを真面目にやってない
• インストーラーを作ってない
• フォーマットや、lint的なことをやってない
• ビルドが重い