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. 成果物は以下

https://github.com/sonesuke/postgresql-ja-tokenizer

# 要件定義

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

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

# 技術選定してみる

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

https://multicorn.org/

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

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

https://www.postgresql.jp/document/13/html/xfunc-c.html

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

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

https://github.com/zombodb/pgx

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

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

# cargo-pgxってなに？

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

<embed>https://www.youtube.com/watch?v=RORkgaURcS0&t=223s</embed>

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

# 用意するもの

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

# セットアップ

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

```bash
$ cargo install cargo-pgx
```

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

```bash
$ cargo pgx init
```

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

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

```bash
$ cargo pgx new ja_tokenizer
```

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

```bash
$ cd ja_tokenizer
$ cargo pgx run pg14
```

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

```sql
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を使用します。

https://github.com/lindera-morphology/lindera

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

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

```rust
#[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で試してみます。

```sql
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をあさってみると、それっぽいのがあるので、以下を参考にやってみます。

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

https://github.com/zombodb/pgx/blob/master/pgx-examples/strings/src/lib.rs

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

```rust
#[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で試してみます。

```sql
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がたくさん使われているので、きっとあってると願いつつ。。。。

https://zenn.dev/frozenlib/articles/lazy_static_to_once_cell

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

```rust
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拡張を作れるようになるまで半日でした。

それでは、良いお年を！

# 積み残し

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

- <del>ユーザ辞書を使いたい</del> (解説記事は[こちら](https://qiita.com/sonesuke/items/0434bccee91b523ba191))
- エラー処理を真面目にやってない
- ドキュメンテーションを真面目にやってない
- インストーラーを作ってない
- フォーマットや、lint的なことをやってない
- ビルドが重い