基於 apache-arrow 的 duckdb rust 客戶端

背景

duckdb 是一個 C++ 編寫的單機版嵌入式分析型數據庫。它剛開源的時候是對標 SQLite 的列存數據庫，並提供與 SQLite 一樣的易用性，編譯成一個頭文件和一個 cpp 文件就可以在程序中使用，甚至提供與 SQLite 兼容的接口，因此受到了很多人的關注。

本文介紹筆者近期開發的 duckdb-rs 庫，讓大家可以很方便地在 rust 代碼庫中使用 duckdb 的功能。

libduckdb-sys

瞭解過 rust 的同學可能知道，rust 提供了 ffi 的方式與其他語言互通。因爲 duckdb 本身是 C++ 編寫的，想要在 rust 裏面使用 duckdb，就需要考慮 ffi 的問題。而基於 ffi 對其他語言程序封裝的基礎庫，一般會被命名爲 libxxx-sys，這也就是 libduckdb-sys 的由來。

爲了方便大家使用，duckdb 提供了 C++ 原生接口，C 接口，以及與 SQLite3 兼容的 C 接口。我在做 libduckdb-sys 的時候對這三種接口都嘗試過，相關的討論可以參見 Rust Support，我這裏介紹一下當時的情況。

基於 SQLite3 接口

最開始我使用的是 SQLite3 的接口，原因主要有三個：

  1. 我對 SQLite 比較熟悉，想必用起來會比較方便；

  2. 覺得 SQLite 的接口被廣泛使用，接口比較穩定，以後不至於大改；

  3. 也許是最重要的一點，市面上已經有 SQLite 的 rust 封裝 rusqlite，基於 SQLite 的接口應該能最大程度複用 rusqlite 的代碼。

嘗試之後確實發現很快能把程序跑起來，基本的功能也能使用。但是隨着進一步的深入以及對 duckdb 更多的瞭解，發現了一些弊端：

  1. 雖說 duckdb 是想最大程度兼容 SQLite，但是畢竟一個是行存一個是列存，有區別在所難免，接口肯定也沒辦法做到 100% 兼容；

  2. 有一個區別需要特別提出來，SQLite 是動態數據類型，而 duckdb 是靜態類型，也就是說在 SQLite 中你可以認爲所有的數據都是存成 Text，在讀取的時候根據 schema 來解析數據；而 duckdb 是會根據數據類型來存儲數據，並且根據列存的特性做一些存儲優化。有了這個區別之後，如果我們使用 SQLite 的接口的話，會做一些不必要的數據格式轉換，性能有損，程序也不直觀。

  3. duckdb 可以被編譯成一個 so 使用，如果想使用 SQLite 的接口，需要再編譯一個 sqlite3_api_wrapper 出來，兩個庫合作才能使用 SQLite 的接口，這給程序分發引入了額外的負擔；另外目前 duckdb 在 release 的時候沒有自帶 sqlite3_api_wrapper，需要用戶自己去編譯，使用上又多了一些不便。

  4. 由於上面的封裝的問題，數據類型的問題，以及通過 SQLite 接口查詢 duckdb 的數據時候，結果集會被複制一遍，資源佔用必定上升。

基於上面一些原因，我最終放棄了基於 SQLite 接口來開發，轉而嘗試使用原生的 C++ 或者 C 接口。

基於 C++ 接口

既然爲了性能和接口豐富性，使用 C++ 接口當然是首選，畢竟 duckdb 本身主要都是拿 C++ 開發的，duckdb 的 python 封裝 也是拿 C++ 接口來做的。

市面上也有方便 rust 與 C++ 交互的一些代碼庫，比如 cxx 和 autocxx。其中 autocxx 入手門檻低使用上更簡單，而 cxx 的可定製性更強，功能更豐富。在嘗試了幾次之後發現了一些問題，主要還是 rust ffi 只能支持部分的 C++ 語法，大部分情況下可能是夠用的，但是對於 duckdb 這樣比較大型的數據庫代碼，還是有很多不支持的地方。除非自己再基於現有的 C++ 接口封裝一份支持 cxx 的版本，否則就算這一次編譯過了，也很難保證以後 duckdb 的作者以後不會引入其他的特性導致不能兼容。

而 rust 基於 C 語言的 ffi 是原生支持的，所以最終還是下定決心基於 C 接口來開發。

基於 C 接口

因爲有 rusqlite 作爲參考，所以很快實現了基於 C 接口的版本。簡單來說，主要是通過 cbindgen、build.rs 和 rust 的 features 功能來實現。其中：

應該說這是一個很通用的提供 C 接口 rust 封裝的解決方案，感興趣的同學可以 參考。

duckdb-rs

完成了 libduckdb-sys 之後其實只是第一步，因爲這樣生成的代碼都是 unsafe 代碼，具體的使用例子可以參考 lib.rs 中的測試代碼。但是我們使用 rust 主要是爲了他的安全性，rust 希望我們儘量減少 unsafe 的使用。所以一般的 rust 封裝都會基於 libxxx-sys 提供一個內存安全的版本，這就是 duckdb-rs 的部分。

小試牛刀

還是因爲有 rusqlite 的參考，所以花了一些時間終於實現了最初始的版本，並且我已經把這個版本發佈到 crates.io 上了。這個版本的目標是基於 rusqlite 做最小的改動，並刪掉 SQLite 特有的功能，讓整個程序跑起來。完成之後效果不錯，下面是文檔中給的一個使用範例：

use duckdb::{params, Connection, Result};

#[derive(Debug)]
struct Person {
    id: i32,
    name: String,
    data: Option<Vec<u8>>,
}

fn main() -> Result<()> {
    let conn = Connection::open_in_memory()?;

    conn.execute_batch(
        r"CREATE SEQUENCE seq;
          CREATE TABLE person (
                  id              INTEGER PRIMARY KEY DEFAULT NEXTVAL('seq'),
                  name            TEXT NOT NULL,
                  data            BLOB
                  );
        ")?;

    let me = Person {
        id: 0,
        name: "Steven".to_string(),
        data: None,
    };
    conn.execute(
        "INSERT INTO person (name, data) VALUES (?, ?)",
        params![me.name, me.data],
    )?;

    let mut stmt = conn.prepare("SELECT id, name, data FROM person")?;
    let person_iter = stmt.query_map([], |row| {
        Ok(Person {
            id: row.get(0)?,
            name: row.get(1)?,
            data: row.get(2)?,
        })
    })?;

    for person in person_iter {
        println!("Found person {:?}", person.unwrap());
    }
    Ok(())
}

可以看到，接口設計非常優雅，代碼也非常符合 rust 的風格，使用上也非常方便。實現過程中發現有些 duckdb 的 C 接口還不支持的部分，我也通過提 issue 或者 PR 去解決了。這裏必須要提一點，duckdb 的維護者非常耐心，不管是回答問題還是 review 代碼都非常專業。

剩下的問題有一個是之前提到的，duckdb 是靜態類型的數據，所以需要支持很多數據類型，這裏面工作量不小。另外，因爲我之前也有關注 Apache Arrow，做過 OLAP 數據庫的同學可能知道，Apache Arrow 是一個通用的列式內存格式，方便在內存中做大數據量的計算或者傳輸，有很多 OLAP 數據引擎都在用。剛好 duckdb 也支持 arrow 格式，所以就想嘗試使用 arrow 格式來查詢數據，這樣至少有兩個好處，一個是這樣我們就可以暴露 arrow 格式的數據給用戶，在使用的時候就可以用上 arrow 生態的其他功能，有可能會產生一些化學反應；另外 arrow 也是有豐富的數據類型和明確的定義，反正我們是要支持很多數據類型的，現在的 C 接口本身也不完善，用 arrow 格式反而更加清晰。

通過 Apache Arrow 查詢數據

基於上面的考慮，我把目標又看向了 arrow-rs，並給 duckdb 的 C 接口也加上了 arrow 的功能，最終在 duckdb-rs 中實現了通過 Arrow 格式來查詢數據，實現參見 這裏。

實現之後，之前通過行來讀取數據的接口完全不變，還能直接查詢到 Arrow 格式的數據，下面是一個測試的例子：

fn test_query_arrow_record_batch_large() -> Result<()> {
    let db = Connection::open_in_memory().unwrap();
    db.execute_batch("BEGIN TRANSACTION")?;
    db.execute_batch("CREATE TABLE test(t INTEGER);")?;
    for _ in 0..300 {
        db.execute_batch("INSERT INTO test VALUES (1); INSERT INTO test VALUES (2); INSERT INTO test VALUES (3); INSERT INTO test VALUES (4); INSERT INTO test VALUES (5);")?;
    }
    db.execute_batch("END TRANSACTION")?;
    let rbs = db.query_arrow("select t from test order by t", [])?;
    assert_eq!(rbs.len(), 2);
    assert_eq!(rbs.iter().map(|rb| rb.num_rows()).sum::<usize>(), 1500);
    assert_eq!(
        rbs.iter()
            .map(|rb| rb
                .column(0)
                .as_any()
                .downcast_ref::<Int32Array>()
                .unwrap()
                .iter()
                .map(|i| i.unwrap())
                .sum::<i32>())
            .sum::<i32>(),
        4500
    );
    Ok(())
}

可以看到，我們查詢到 Arrow 格式的數據之後，還能通過 arrow-rs 中提供的其他能力做進一步的計算，十分方便。

總結

本文主要介紹了 duckdb-rs 的設計和實現，筆者之前有一些開發 OLAP 數據的經驗，但是對於 rust 算是新手，之前雖然寫過一些但是沒有深入學習，做這個項目也有一個目的是爲了重新學習一下 rust。好在有 rusqlite 作爲參考，所以沒有碰到特別多語言層面的問題。

希望這篇文章對於其他對 rust 和數據庫感興趣的同學有一些幫助。同時這個庫還有很多沒解決的問題，比如支持更多的數據類型，支持連接池，支持更快的數據導入接口等等，我已經建了一些 issues，感興趣的同學可以回覆 issue 認領，我也會竭力提供需要的幫助，大家一起討論和學習。

參考

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/lSl-3rMNufsDxhNuAbZ93Q

猜你喜歡