DuckDB の拡張を Rust で書けるらしいので調べてみた（2025年3月時点）

DuckDB は v1.2.0 で C extension API というものが追加されました。これによって、C で拡張を書けるのみならず、Rust や Go などの C FFI を使えるプログラミング言語でも（その C extension API へのバインディングがあれば）書けるようになりました。

https://duckdb.org/2025/02/05/announcing-duckdb-120.html

実際、Rust 実装の拡張を書くテンプレートが公式に用意されています。

https://github.com/duckdb/extension-template-rs

とはいえ、このテンプレートはまだ「experimental」だと書かれています。現状、これを使って Rust で DuckDB の拡張を書くのはどれくらいできるもんなの？、というのを調べてみました。

拡張の種類

DuckDB の拡張でできることの範囲が分かっていないのですが、API 一覧を見たところ、拡張で提供できる機能には以下があるようです。

Table function

ドキュメント: https://duckdb.org/docs/stable/clients/c/table_functions

Table function は、テーブルを提供する関数です。要は、FROM のあとに置けるものです。read_parquet() とか read_excel() とか、そういうやつです。

SELECT * FROM fun();

ちなみに、副作用のためだけに使う関数（CALL fun() のように使うやつ）も Table function として実装すればいいみたいです。 start_ui() の実装を見ていて気付きました。

Salar function

ドキュメント: （なし）

Scalar function は、値を受け取って同じ数の値を返す関数です。要は、fun(<列名>) という使い方をできるものです。toupper() とか sqrt() とか、そういうやつです。

SELECT fun(col1) FROM t1;

Aggregate function

ドキュメント: （なし）

Aggregate function は、集約関数（複数の値を受け取って1つの値を返す関数）です。要は、GROUP BY に使えるものです。sum() とか all() とかそういうやつです。

SELECT fun(col1) FROM t1 GROUP BY col2;

Replacement scan

ドキュメント: https://duckdb.org/docs/stable/clients/c/replacement_scans

ここまでは関数が並んでいましたが、これはちょっと特殊なものです。ドキュメントには「SELECT * FROM my_table で my_table が見つからなかったときに呼ばれる処理」といった感じで書かれていますが、わかりやすく言えば、

SELECT * FROM 'path/to/data.parquet'

というクエリがあったときに、自動で read_parquet('path/to/data.parquet') を呼んでくれる、といった仕組みです。拡張で Table function を提供するときには、Replacement scan も実装しておくと快適に使えそうです。

duckdb-rs の実装状況

さて、上には4つの機能を挙げましたが、duckdb-rs で提供されているのは、今のところ以下の2つだけです。

• Table function
• Scalar function

しかも、Scalar function の方は、先月追加されたばかり（duckdb/duckdb-rs#429）で、まだ feature flag をうまく選ばないとコンパイルできなかったり、いろいろ改善の余地がありそうな状態のまま止まっています。

実は、duckdb-rs は最近はあまりアクティブには開発されていません（参考：コミット履歴）。たぶん、DuckDB の中で Rust わかる人が優先度的に他のところをやっていて、まだ取り掛かれていないのかな...？と思ったりしています。

そんなわけで、Rust で DuckDB 拡張を書くのはまだ早い、というのが私の中での結論です。
まあでも、ふつうに動くものは作れるので、以下では具体的なコードについて軽く説明していきます。

Table function を実装するには

Table function を実装するのは、duckdb-rs の example に例があります。

https://github.com/duckdb/duckdb-rs/blob/1.2.0/crates/duckdb/examples/hello-ext/main.rs

具体的には、VTab trait、もしくは Arrow 版である VArrowTab trait を実装して、register_table_function() で関数を登録します。Arrow の API は deprecated であるという情報もあるので^[1]、ここでは VTab の方を見ていきましょう。

pub trait VTab: Sized {
    type InitData: Sized + Send + Sync;
    type BindData: Sized + Send + Sync;

    // Required methods
    fn bind(bind: &BindInfo) -> Result<Self::BindData, Box<dyn Error>>;
    fn init(init: &InitInfo) -> Result<Self::InitData, Box<dyn Error>>;
    fn func(
        func: &TableFunctionInfo<Self>,
        output: &mut DataChunkHandle,
    ) -> Result<(), Box<dyn Error>>;

    // Provided methods
    fn supports_pushdown() -> bool { ... }
    fn parameters() -> Option<Vec<LogicalTypeHandle>> { ... }
    fn named_parameters() -> Option<Vec<(String, LogicalTypeHandle)>> { ... }
}

bind()

bind というのがどういう概念なのかよくわかっていないのですが、どうやらここは

• 結果のカラムの名前と型を設定（add_result_column()）
• 予想される結果の行数を設定（set_cardinality()）

といった、結果のメタデータの設定を行うステージのようです。

また、ユーザーが table funciton に指定した引数を受け取るのもここでやります。引数として入ってくる BindInfo に get_parameter() というメソッドがあるので、それで値を取り出して、必要な情報をまとめて Self::BindData（associated type なので自分で好きに定義できる） として return します。この Self::BindData は、あとで func() の中で参照できるので、そこでパラメータによって処理を切り替えたりすることができます。

init()

init は bind のあとに実行されるステージで、ファイルをオープンしたりコネクションを作ったりするところのようです。bind との使い分けがよくわからないのですが、状態を保持するものは Self::BindData ではなく Self::InitData の方にここで入れる、というのがお作法のようです。

実際にファイルやコネクションを扱っている例は Rust のコードではまだなさそうなので、正しいやり方はいまいちよくわかりません。

func()

ここがメインの処理です。引数として入ってくる output にデータを詰め込みます。
DataChunkHundle の扱い方については、長くなるので下の「DataChunkHundle にデータを書くには」で説明しますが、data chunk という概念についてだけ触れておくと、ドキュメントには以下のように説明されています。テーブルを横にぶつ切りにしたもの、みたいなイメージで私は理解しています。要は、DuckDB がデータを小分けにして処理するその単位です。

Data chunks represent a horizontal slice of a table. They hold a number of vectors, that can each hold up to the VECTOR_SIZE rows. The vector size can be obtained through the duckdb_vector_size function and is configurable, but is usually set to 2048.

この処理は並列で実行されうるので、Self::InitData で管理する変数（seek の位置とか）は Mutex などでラップしておく必要があります。たとえば、↑の example では、AtomicBool を使うことで 1 度しか読まれないことを担保しています。

parameters(), named_parameters()

table function に引数を付ける場合はここで設定します。positional な引数なら parameters()、named な引数なら named_parameters() です。

引数の型は LogicalTypeId という enum から選べばいいだけですが、指定するのは LogicalTypeId ではなくて LogicalTypeHandle なので、from() とか into() で変換が必要です。

fn parameters() -> Option<Vec<LogicalTypeHandle>> {
    Some(vec![LogicalTypeId::Varchar.into()])
}

supports_pushdown()

これは使い方がわかりません...

Scalar function を実装するには

Scalar function は、example には例がないんですが、テストのコードが参考になります。

https://github.com/duckdb/duckdb-rs/blob/1096461ea9d6a9632e5cb079860ea47b2de0257b/crates/duckdb/src/vscalar/mod.rs#L196-L238

これも Table function と同じく、通常版の trait（VScalar）と Arrow 版の trait（VArrowScalar） があります。ここでは VScalar を見ていきましょう。

pub trait VScalar: Sized {
    type State: Default + Sized + Send + Sync;

    // Required methods
    unsafe fn invoke(
        state: &Self::State,
        input: &mut DataChunkHandle,
        output: &mut dyn WritableVector,
    ) -> Result<(), Box<dyn Error>>;
    fn signatures() -> Vec<ScalarFunctionSignature>;
}

ちなみに、おそらくですけど、State はこの trait を実装する struct 自身が兼ねればよくて、こんな感じでいいのでは...？と思っています。associated type として持っておく利点が私にはよくわかっていません（単純に何か見落としてるだけなのかもです）。

pub trait VScalar: Default + Sized + Send + Sync {
    // Required methods
    unsafe fn invoke(
        &mut self,
        input: &mut DataChunkHandle,
        output: &mut dyn WritableVector,
    ) -> Result<(), Box<dyn Error>>;
    fn signatures() -> Vec<ScalarFunctionSignature>;
}

invoke()

ここに実際の処理が入ります。input に入力が入ってくるので、それに対して何らかの処理を行って output に書き込みます。入力の読み取りは、下の「DataChunkHundle からデータを読むには」で説明します。output の WritableVector の扱いも下の「DataChunkHundle にデータを書くには」で説明します。

scalar funciton に何らかの状態を持たせたい場合は、Self::State に入れます。state は mutable ではないので、実際の変数は Mutex などでラップすることになると思います。特に状態が必要なければ、() を入れておけばいいみたいです。

impl VScalar for MyScalar {
    type State = ();

signatures()

関数のシグネチャを返します。戻り値がベクトルになっているのは、scalar function は複数のシグネチャを持つことができるからです。つまり、入力が integer なら integer で返して、入力が float なら float で返して、みたいなことです。

引数は、可変長にすることもできます。固定長であれば ScalarFunctionSignature::exact() 、可変長引数なら ScalarFunctionSignature::variadic() を使います。可変長引数の場合は、今のところ、すべての引数で型が同じやつしか許されないようです（例えば foo(col1, col2, ..., option) みたいなことはできなそう）。

たとえば、入力が integer の引数1つ、出力も integer、という関数のシグネチャはこうなります。

    fn signatures() -> Vec<duckdb::vscalar::ScalarFunctionSignature> {
        vec![ScalarFunctionSignature::exact(
            vec![LogicalTypeId::Integer.into()],
            LogicalTypeId::Integer.into(),
        )]
    }

このシグネチャに合わない型の引数が指定されると、クエリ実行時に以下のようなエラーになります。つまり、invoke() の方の処理は、ここで指定した型の入力しか入ってこないという想定で実装してよさそうです。ただし、シグネチャが複数あるならその分の場合分けは必要になります（後述）。

Binder Error:
No function matches the given name and argument types 'hello_scalar(DECIMAL(2,1))'. You might need to add explicit type casts.
        Candidate functions:
        hello_scalar(INTEGER) -> INTEGER

DataChunkHundle からデータを読むには

流れとしてはこのようになります。

// 入力の行数を調べる
let len = input.len();

// 1列目を flat vector として取り出す
let input_vec = input.flat_vector(0);

// Rust で読める slice に変換（入力が i32 のみだと仮定）
let input_values = input_vec.as_slice_with_len::<i32>(len);

• 可変長引数の場合は、input.num_columns() で列数も調べる必要があるかもです。
• flat vector というのは、primitive type のみのベクトルです。他にも、list_vector()、array_vector()、struct_vector() などがあるようです。
• FlatVector には as_slice() と as_slice_with_len() という2種類のメソッドが生えていますが、as_slice() だと len() じゃなくて capacity() まで取ってしまうので、as_slice_with_len() が正解のようです（as_slice() を使う時はあるんだろうか...？）。

注意点として、上の例では入力が i32 のみだと仮定しているので簡単なコードになっていましたが、シグネチャが複数ある場合は LogicalTypeId によって場合分けをする必要があります。イメージを示すと、以下のような感じです。ただし、実際にはこのコードは &[i32] と &[f32] で型が違ってコンパイルできません。型を合わせるように書くのがちょっと面倒そうです（ジェネリクスを作って別の関数に切り出すことになる...？）。

let input_values = match input_vec.logical_type().id() {
    LogicalTypeId::Integer => input_vec.as_slice_with_len::<i32>(len),
    LogicalTypeId::Float => input_vec.as_slice_with_len::<f32>(len),
    ...
    id => {
        return Err(format!("Unsupported type: {id:?}").into());
    }
};

flat vector の中でも、文字列や Blob の場合はポインタが入っている点に注意しましょう。DuckString::new() を使うと実際の値が取り出せます。

let strings = values
    .iter()
    .map(|ptr| DuckString::new(&mut { *ptr }).as_str().to_string());

DataChunkHundle にデータを書くには

基本的な流れは、

1. flat vector として列を取り出す
2. as_mut_slice_with_len() とかで &mut [T] にしてそこに値を書き込んでいく

という感じになるみたいです。

let flat_vec: &mut [i32] = output.flat_vector(0)
    .as_mut_slice_with_len::<i32>(len);

文字列や Blob の場合は Inserter という trait が用意されていて、insert() が使えます。

https://docs.rs/duckdb/latest/duckdb/core/trait.Inserter.html

https://github.com/duckdb/duckdb-rs/blob/6d5368011043bd65dac42cf5c7a26784c16d82cc/crates/duckdb/src/vscalar/mod.rs#L213-L230

WritableVector

VScalar::invoke() に出てくる WritableVector は、扱いとしては DataChunkHundle の1列バージョンみたいなものだと考えればよさそうです。結局 Flat vector を取り出してそこに書き込むのは同じで、DataChunkHundle では .flat_vector(列番号) でデータを取り出すところが、.flat_vector() と引数なしになるだけです。

具体的にはこんな感じです（入力の i32 を2倍にして返す例）。

let mut flat_vec = output.flat_vector();
flat_vec
    .as_mut_slice_with_len::<i32>(len)
    .iter_mut()
    .zip(input_values.iter())
    .for_each(|(o, i)| {
        *o = 2 * i;
    });

余談

実は、せっかく記事を書くから、Aggregate function を duckdb-rs に実装するぞ！と意気込んでいたのですが、想像以上に大変で挫折しました。Scalar function のコピペでいけるだろう、と思っていたら、その10倍くらい複雑でした。これがその残骸です：

https://github.com/duckdb/duckdb-rs/compare/1096461...yutannihilation:duckdb-rs:feat-aggregate-function

というのは、Scalar function は入ってきた値を1:1で返すだけなので複雑な状態を持つ必要はないのですが、Aggregate function はそういうわけにはいきません。途中の計算結果をちゃんと状態として持っておかないといけなくて、そのためには DuckDB 側がどういう処理の中でこの API を呼び出すのかを知っておかないとうまく実装できなそうでした。とりあえず私は、どういうインターフェースなのか理解が及ばず諦めました...

これが単に私の理解力の問題なのか、それとも DuckDB の API がまだ整理されていない状態だからなのかは正直よくわかりません。とりあえずやってみてわかったのは、拡張のために適切に API を切るというのは難易度の高いことであり、いろいろ試行錯誤しないとわからないこともありそう、ということでした。なので、C extension API ができても、すぐに duckdb-rs が実装されないのはまあ理解できるというか、いろいろ考えることあるんだろうな、と思ったりしました。

今後に期待！

脚注

1. とはいえ、現状の duckdb-rs は Arrow にかなり依存しているので、すぐなくなることはないと思います。DuckDB 本体側としても、非推奨だからと言って移行先があるわけではないという状態のようなので（参考）、Arrow 関連の API を実際に削除できるのかは不明です。