DuckDB の拡張を Rust で書く方法を調べる
yutannihilation記事にまとめたのでこのスクラップはクローズします。
yutannihilationDuckDB は v1.2.0 で C extension API というものが追加された。
これを使って拡張を開発する利点は以下。
- DuckDB を静的リンクしなくてもいいので軽くなるし、DuckDB のバージョンが変わっても動く（C API のバージョンが変われば互換性が壊れることはあり得る）
- C ABI のバイナリが生成できる言語（Rust とか Go とか）で開発できる
ということで、Rust の DuckDB 拡張公式テンプレートが用意されている。
が、まだいろいろ整ってなくていま触るのはいばらの道っぽい...、というのをメモっていく。ある程度まとまったらスクラップじゃなくて記事として書き直したい。
yutannihilationDuckDB の Rust binding としては duckdb-rs がある。
これは、C API（DuckDB 自体を Rust から操作したりするやつ）も C extension API も両方カバーしていて、デフォルトだと C API 用になっている。 loadable-extension という feature を有効にすると、C extension API に切り替わる。具体的には、C API と C extension API ではヘッダファイルが違っていて（API として用意されているものが違うので当然）、このへんで切り替わっている。
ちなみに、細かいことを言えば、「libduckdb-sys」は sys crate だが、上述のように C extension API の場合は DuckDB に直接リンクしないはずなので -sys でなはないのが正しい気がする。まあでもたぶん実害はないし、別に crate つくるのはめんどくさいし、これでいいんだとは思う。
デフォルトが C API なので、docs.rs には C extension API 用の関数は含まれていない。とりあえず自分でビルドしたものを以下に置いてみた（メンテするつもりはないのでそのうち消す）
yutannihilationDuckDB 拡張でできる範囲がよくわかっていないが、とりあえず、 「register」という単語で duckdb_extension.h を検索した感じ、以下の 3 種類があるっぽい。
- table function: テーブルを返す関数。
- scalar function: 値を返す関数。
- aggregate function: 集約関数。
このうち、extension-template-rs で例として書かれているのは table function のみ。というのは、少し前までは table function しかなかったからっぽい。scalar function が実装できるようになったのが2月11日（pull request）、aggregate function はまだ使えない。
さらに、scalar function には通常版、Arrow 版の2種類があるらしい。arrow の方が速いけど、arrow-rs が必要になるのでビルドがちょっと重くなる、ということなんだと思う。
たぶん aggregate function も同じく通常版と Arrow 版ができることになりそう。table function は ArrowVTab っていうのがあるけど、これは trait じゃなくて struct なので関係ない。
ちょっとわからなかったのは、VScalar は invoke() というメソッドを実装しないといけないけど、そのシグネチャがこうなっていて、 WritableVector が Arrow のものなので、結局 Arrow ないと使えない...？という予感がしている。とりあえず issue 立ててみた → duckdb/duckdb-rs#441
unsafe fn invoke(
state: &Self::State,
input: &mut DataChunkHandle,
output: &mut dyn WritableVector,
) -> Result<(), Box<dyn Error>>;
とりあえずここでいろいろ実験しています。
yutannihilationScalar function を実装する
上に書いたように、scalar function 用の trait としては、VScalar と VArrowScalar がある。
どちらもメソッドは2つ。input を受け取って output を返す（つまり実際の処理を書く） invoke() と、シグネチャを定義する signatures()。違いは、 invoke() が output を受け取ってそこに書き込むか、戻り値（Arc<dyn arrow::array::Array>）として返すかという点だけ。
pub trait VScalar: Sized {
type State: Default + Sized + Send + Sync;
unsafe fn invoke(
state: &Self::State,
input: &mut DataChunkHandle,
output: &mut dyn WritableVector,
) -> Result<(), Box<dyn Error>>;
fn signatures() -> Vec<ScalarFunctionSignature>;
}
pub trait VArrowScalar: Sized {
type State: Default + Sized + Send + Sync;
fn invoke(
info: &Self::State,
input: RecordBatch,
) -> Result<Arc<dyn Array>, Box<dyn Error>>;
fn signatures() -> Vec<ArrowFunctionSignature>;
}
signatures()
まずは簡単なこっちから。こういう感じで書くといいらしい。これは、入力が 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
整数も実数も OK にしたいんだけど、という場合のために Vec になっているので、好きなだけシグネチャを並べればいい。ただし、シグネチャが1つでない場合は invoke() の方で場合分けが必要になったりしてちょっとコードがややこしくなりそう。
ScalarFunctionSignature::exact() は固定長引数の関数。可変長引数がいい場合は ScalarFunctionSignature::variadic() というのが用意されている。可変長引数の場合は、今のところ、すべての引数で型が同じやつしか許されない（例えば foo(col1, col2, ..., option) みたいなことはできなそう）っぽい。
VScalar の方は、型の指定が LogicalTypeId ではなくて LogicalTypeHandle なので into() で変換している。VArrowScalar だと DataType が指定できる。
invoke()
input
入力は DataChunkHandle に入ってくる。Data Chunk とは何かというと、
Data chunks represent a horizontal slice of a table. They hold a number of vectors, that can each hold up to the
VECTOR_SIZErows. The vector size can be obtained through theduckdb_vector_sizefunction and is configurable, but is usually set to 2048.
ということで、テーブルを横にぶつ切りにしたものだと理解した。ここから各列を取り出すには、こんな感じでやるらしい。
let len = input.len();
let input_vec = input.flat_vector(0);
let input_values = input_vec.as_slice_with_len::<i32>(len);
- 1行目： まずは行数を取る。ちなみに列数は
num_columnsでわかるらしい。 - 2行目： 1列目（index が 0）を flat vector（primitive type のみのベクトル？）として取り出している。他にも
list_vector()、array_vector()、struct_vector()などがあるらしいが、使ってる実例がなくてわからなかった。 - 3行目：
FlatVectorにはas_slice()とas_slice_with_len()という2種類のメソッドが生えているが、例ではas_slice_with_len()しか使われていなかった。なんでかなと思ってコードを見たら、as_slice()だとlen()じゃなくてcapacity()まで取ってしまうので、たしかにas_slice_with_len()が適切っぽい（as_slice()を使う時あるんだろうか...？）。
そして、as_slice_with_len::<T>() の T は、シグネチャが1つの場合は↑のように特定の値を書けばいいが、複数の場合は 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());
}
};
文字列や Blob の場合はポインタが入っている。DuckString::new() を使うと実際の値が取り出せるらしい。
let strings = values
.iter()
.map(|ptr| DuckString::new(&mut { *ptr }).as_str().to_string());
output
基本的な流れは、
- flat vector として取り出す
-
as_mut_slice_with_len()とかで&mut [T]にしてそこに値を書き込んでいく
という感じになりそう。具体的にはこんな感じ（入力の 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;
});
文字列や Blob の場合は Inserter という trait が用意されていて、insert() が使える。
yutannihilationC API のうち、Arrow interface はどうやら deprecated らしい。deprecated だからといって移行先があるわけではないし、廃止する計画があるわけでもないらしいので使い続けられることになるのだと思う。duckdb-rs はこれに依存してるのでちょっと不安。