SurrealDBを組み込み実行してJSONを検索

SurrealDB を組み込み実行して、「DuckDBを使ってJSONをSQLで処理」と同等の処理を実施してみました。

組み込み実行

Surreal::new で surrealdb::engine::local::Mem の型を指定するとインメモリで組み込み実行できます。

実装例は次のようになり、use_ns でネームスペースを use_db でDB名を指定します。

src/main.rs

use surrealdb::Surreal;
use surrealdb::engine::local::Mem;

...省略

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

#[tokio::main]
async fn main() -> Result<()> {
    let db = Surreal::new::<Mem>(()).await?;
    db.use_ns("sample").use_db("db").await?;

    ...省略

    Ok(())
}

Cargo.toml はこのようにしました。

Cargo.toml

[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
surrealdb = { version = "2.3.6", features = ["kv-mem"] }
tokio = { version = "1", features = ["macros", "rt-multi-thread"] }

データ登録

まずは、下記ファイルの JSON をそれぞれ登録しておきます。

data/items.jsonl

{"id":1,"name":"item1","price":1100,"attrs":{"code":"0001","category":"A1"},"variants":[{"color":"white", "size":"S"},{"color":"black","size":"F"}]}
{"id":2,"name":"item2","price":2300,"attrs":{"code":"0002"},"variants":[{"color":"red"},{"color":"blue"},{"color":"white"}]}
{"id":3,"name":"item3","price":360,"attrs":{"code":"0003"}}
{"id":4,"name":"item4","price":4700,"attrs":{"code":"0004","category":"A1"},"variants":[{"color":"green", "size":"S"}]}
{"id":5,"name":"item5","price":590,"attrs":{"code":"0005","category":"B2"},"variants":[{"color":"red"},{"color":"white","size":"F"},{"color":"white","size":"L"}]}
{"id":6,"name":"item6","price":6200,"attrs":{"code":"0006","category":"C3"},"variants":[{"color":"blue","size":"S"},{"color":"green","size":"L"}]}

SurrealDB では create(テーブル名) もしくは create((テーブル名, id)) でリソース（レコード）を作成して、content で内容を設定できます。

任意の id を指定する方法として、次の 2通りあるようですが¹、ここでは 2つ目の方法を使っています。

• create の引数をタプルにして id を指定
• content で設定する内容で id を指定

任意の型²を使って content へ設定できますが、ここでは serde_json::Value を使いました。

また、戻り値の型を明示的に指定する必要がありますが、こちらは定義した型（Document）を使ってみました。

データ登録

...省略
use surrealdb::sql::Thing;

use serde::Deserialize;
use serde_json::Value;

use std::fs::File;
use std::io::{BufRead, BufReader};
// create で取得する型
#[allow(dead_code)]
#[derive(Debug, Deserialize)]
struct Document {
    id: Thing,
}

...省略

#[tokio::main]
async fn main() -> Result<()> {
    let db = Surreal::new::<Mem>(()).await?;
    db.use_ns("sample").use_db("db").await?;

    let file = "./data/items.jsonl";

    for line in BufReader::new(File::open(file)?).lines() {
        if let Ok(s) = line {
            let v: Value = serde_json::from_str(&s)?;
            // レコードの作成
            let r: Option<Document> = db.create("items").content(v).await?;

            println!("* registered: {:?}", r);
        }
    }
    ...省略
}

実行結果はこのようになります。

実行結果

* registered: Some(Document { id: Thing { tb: "items", id: Number(1) } })
* registered: Some(Document { id: Thing { tb: "items", id: Number(2) } })
* registered: Some(Document { id: Thing { tb: "items", id: Number(3) } })
* registered: Some(Document { id: Thing { tb: "items", id: Number(4) } })
* registered: Some(Document { id: Thing { tb: "items", id: Number(5) } })
* registered: Some(Document { id: Thing { tb: "items", id: Number(6) } })

クエリ処理

SurrealDB では SQL ライクなクエリ言語を使用できます。

(a) attrs, variants の出力

id が 4未満の id, attrs.code, attrs, variants を検索するクエリはこのようになります。

クエリ例

SELECT id, attrs.code as code, attrs, variants FROM items WHERE id < items:4

id は テーブル名:値 で指定する必要があり、id < 4 では意図した結果を取得できませんでした。³

また、attrs.code だと結果が入れ子になるため、フラットな値で取得するには as で別名を付ける必要がありました。

クエリは query で実行でき、query を複数繋げる事が可能です。
take でどのクエリの結果を取得するかを指定し、例えば 1番目のクエリの結果を取得するには take(0) とします。

ここでも取得するデータの型を明示的に指定する必要がありますが、型定義は行わずに surrealdb::Value を使いました。

クエリ処理

let a = "SELECT id, attrs.code as code, attrs, variants FROM items WHERE id < items:4";
let b = "SELECT ...省略";
let c = "SELECT ...省略";
// a, b, c の 3つのクエリを実行
let mut res = db.query(a).query(b).query(c).await?;
// a の結果を取得
let a_res: surrealdb::Value = res.take(0)?;
println!("(a) = {}", a_res.to_string());

...省略

実行結果はこうなります。

(a)の実行結果

(a) = [{ attrs: { category: 'A1', code: '0001' }, code: '0001', id: items:1, variants: [{ color: 'white', size: 'S' }, { color: 'black', size: 'F' }] }, { attrs: { code: '0002' }, code: '0002', id: items:2, variants: [{ color: 'red' }, { color: 'blue' }, { color: 'white' }] }, { attrs: { code: '0003' }, code: '0003', id: items:3, variants: NONE }]

ちなみに、take する型は以下のようなものが使えるようです。

• surrealdb::Value
• Vec<T>
• Option<T>

例えば、定義した Document 型で取得する場合は res.take::<Vec<Document>>(0) とします。

また、タプルを使って指定の項目だけを取得でき、例えば res.take::<Vec<Thing>>((0, "id")) とすると id だけを取得できます。

(b) attrs.category の条件指定

attrs の category が 'A1' のものだけを抽出するクエリはこのようになります。

(b)のクエリ処理

let b = "SELECT id, name, attrs.category FROM items WHERE attrs.category = 'A1'";

...省略

println!("(b) = {}", res.take::<surrealdb::Value>(1)?.to_string());

(b)の実行結果

(b) = [{ attrs: { category: 'A1' }, id: items:1, name: 'item1' }, { attrs: { category: 'A1' }, id: items:4, name: 'item4' }]

(c) variants の条件指定

配列項目[WHERE 条件] で配列内の要素を条件指定できるため、variants 内に color='white' を含むものを抽出するクエリはこのようになります。

(c)のクエリ処理

let c = "SELECT id, name FROM items WHERE variants[WHERE color = 'white']";

...省略

println!("(c) = {}", res.take::<surrealdb::Value>(2)?.to_string());

(c)の実行結果

(c) = [{ id: items:1, name: 'item1' }, { id: items:2, name: 'item2' }, { id: items:5, name: 'item5' }]

1. これ以外ではランダムな値が適用される様子 ↩

2. serde::Serialize 等の derive が必要 ↩

3. WHERE id < 4 とすると結果は空になり、WHERE id > 4 とすると全レコードがヒットしました ↩