Alopex DB v0.8.8チュートリアル 第4回：ベクトル検索とDataFrame

pip install --no-cache-dir "alopex[polars]==0.8.8"

from alopex import Database

db = Database.new()
db.execute_sql("""CREATE TABLE articles (
  id INTEGER PRIMARY KEY,
  title TEXT,
  tags TEXT,
  views INTEGER,
  embedding VECTOR(3, L2)
)""")

rows = [
    "(1, 'Rustで書くLSMツリー', 'rust,storage', 1200, [1.0, 0.0, 0.0])",
    "(2, 'SQLパーサをNimで実装する', 'nim,sql', 800, [0.9, 0.2, 0.0])",
    "(3, 'ベクトル検索の基礎', 'vector,search', 1500, [0.0, 1.0, 0.0])",
    "(4, '組み込みDBの選び方', 'database,embedded', 600, [0.1, 0.9, 0.0])",
    "(5, 'WALとクラッシュ復旧', 'storage,recovery', 300, [0.8, 0.0, 0.3])",
]
for row in rows:
    db.execute_sql(f"INSERT INTO articles VALUES {row}")

print(db.execute_sql("""
SELECT id, title, vector_distance(embedding, [1.0, 0.0, 0.0], 'l2') AS d
FROM articles
ORDER BY d
LIMIT 3
"""))

[{'id': 1, 'title': 'Rustで書くLSMツリー', 'd': 0.0},
 {'id': 2, 'title': 'SQLパーサをNimで実装する', 'd': 0.22360680997371674},
 {'id': 5, 'title': 'WALとクラッシュ復旧', 'd': 0.36055511236190796}]

print(db.execute_sql("""
SELECT id, title, vector_distance(embedding, [1.0, 0.0, 0.0], 'l2') AS d
FROM articles
WHERE views >= 800
ORDER BY d
LIMIT 3
"""))

[{'id': 1, 'title': 'Rustで書くLSMツリー', 'd': 0.0},
 {'id': 2, 'title': 'SQLパーサをNimで実装する', 'd': 0.22360680997371674},
 {'id': 3, 'title': 'ベクトル検索の基礎', 'd': 1.4142135381698608}]

from alopex import DataFrame

rows = db.execute_sql("SELECT id, title, tags, views FROM articles ORDER BY id")
df = DataFrame.from_columns({k: [r[k] for r in rows] for k in rows[0]})

print(df.height(), "行", df.width(), "列")

5 行 4 列

tagged = df.str("tags").split(",", "tag_list")
print(tagged.to_dict()["tag_list"])

[['rust', 'storage'], ['nim', 'sql'], ['vector', 'search'],
 ['database', 'embedded'], ['storage', 'recovery']]

per_tag = tagged.explode("tag_list")
d = per_tag.to_dict()
print(d["tag_list"])
print(d["id"])

['rust', 'storage', 'nim', 'sql', 'vector', 'search', 'database', 'embedded', 'storage', 'recovery']
[1, 1, 2, 2, 3, 3, 4, 4, 5, 5]

from alopex import col, lit

hot = (per_tag.lazy()
       .filter(col("views").gt(lit(800)))
       .select([col("tag_list"), col("views")])
       .collect())

print(hot.to_dict())

{'tag_list': ['rust', 'storage', 'vector', 'search'], 'views': [1200, 1200, 1500, 1500]}

df = DataFrame.from_columns(
    {"published": [1768473000000000]},
    {"published": "datetime"},
)

[dependencies]
alopex-embedded = "=0.8.8"
alopex-dataframe = "0.8"
alopex-sql = "0.8"
arrow = "53"

use alopex_dataframe::{col, lit, DataFrame, Series};
use alopex_embedded::{Database, SqlResult};
use alopex_sql::storage::value::SqlValue;
use arrow::array::{ArrayRef, Int64Array, StringArray};
use arrow::util::display::array_value_to_string;
use std::path::Path;
use std::sync::Arc;

let result = db.execute_sql(
    "SELECT id, title, tags, views, vector_distance(embedding, [1.0, 0.0, 0.0], 'l2') AS d
     FROM articles ORDER BY d LIMIT 3")?;

let SqlResult::Query(q) = result else { return Ok(()) };
for row in &q.rows {
    println!("{}", row.iter().map(|v| format!("{v:?}")).collect::<Vec<_>>().join(" | "));
}

Integer(1) | Text("Rustで書くLSMツリー") | Text("rust,storage") | Integer(1200) | Double(0.0)
Integer(2) | Text("SQLパーサをNimで実装する") | Text("nim,sql") | Integer(800) | Double(0.22360680997371674)
Integer(5) | Text("WALとクラッシュ復旧") | Text("storage,recovery") | Integer(300) | Double(0.36055511236190796)

let mut ids = Vec::new();
let mut tags = Vec::new();
let mut views = Vec::new();
for row in &q.rows {
    if let SqlValue::Integer(v) = &row[0] { ids.push(*v as i64) }
    if let SqlValue::Text(v) = &row[2] { tags.push(v.clone()) }
    if let SqlValue::Integer(v) = &row[3] { views.push(*v as i64) }
}

let id_col: ArrayRef = Arc::new(Int64Array::from(ids));
let tag_col: ArrayRef = Arc::new(StringArray::from(tags));
let view_col: ArrayRef = Arc::new(Int64Array::from(views));

let df = DataFrame::new(vec![
    Series::from_arrow("id", vec![id_col])?,
    Series::from_arrow("tags", vec![tag_col])?,
    Series::from_arrow("views", vec![view_col])?,
])?;
println!("{} 行 {} 列", df.height(), df.width());

3 行 3 列

let per_tag = df.lazy()
    .with_columns(vec![col("tags").str().split(",").alias("tag_list")])
    .explode("tag_list")
    .collect()?;
println!("--- explode: {} 行", per_tag.height());

--- explode: 6 行

let hot = per_tag.lazy()
    .filter(col("views").gt(lit(700i64)))
    .select(vec![col("id"), col("tag_list"), col("views")])
    .collect()?;

for s in hot.columns() {
    let arr = &s.to_arrow()[0];
    let cells: Vec<String> = (0..arr.len())
        .map(|i| array_value_to_string(arr, i).unwrap())
        .collect();
    println!("{}: {}", s.name(), cells.join(", "));
}

id: 1, 1, 2, 2
tag_list: rust, storage, nim, sql
views: 1200, 1200, 800, 800