Rustでcsvバイナリデータを読み込んでsqliteにデータを挿入する

Rust polarsでcsvデータを読み込んで
sqliteに格納するコードの備忘録です。

use polars::prelude::*;
use rusqlite::{params, Connection, Result};
use std::fs::File;
use std::io::Cursor;

fn main() -> Result<()> {
    let conn = Connection::open("test.db")?;

    conn.execute(
        "CREATE TABLE IF NOT EXISTS csv_data (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            data BLOB NOT NULL
        )",
        [],
    )?;

    let csv_file_path = "/Users/ods/Documents/surreal/data/data.csv";
    let df = read_csv_as_dataframe(csv_file_path).expect("CSVファイルの読み込みに失敗しました");

    save_dataframe_as_binary(&conn, "key-name", df)?;

    println!("DataFrameをバイナリ形式でSQLiteに保存しました。");

    let restored_df = restore_dataframe_from_binary(&conn,"key-name")
        .expect("SQLiteからDataFrameの復元に失敗しました");
    println!("{:?}", restored_df);

    Ok(())
}

fn read_csv_as_dataframe(file_path: &str) -> PolarsResult<DataFrame> {
    CsvReader::new(File::open(file_path)?)
        .finish()
}

fn save_dataframe_as_binary(conn: &Connection, name: &str, mut df: DataFrame) -> Result<()> {
    let mut buffer = Vec::new();
    CsvWriter::new(&mut buffer)
        .finish(&mut df)
        .expect("DataFrameのバイナリ変換に失敗しました");

    conn.execute(
        "INSERT INTO csv_data (name, data) VALUES (?1, ?2)",
        params![name, buffer],
    )?;

    Ok(())
}

fn restore_dataframe_from_binary(conn: &Connection, name: &str) -> PolarsResult<DataFrame> {
    let mut stmt = conn
        .prepare("SELECT data FROM csv_data WHERE name = ?1")
        .map_err(|e| PolarsError::ComputeError(format!("SQLiteエラー: {:?}", e).into()))?;
    let binary_data: Vec<u8> = stmt
        .query_row(params![name], |row| row.get(0))
        .map_err(|e| PolarsError::ComputeError(format!("SQLiteエラー: {:?}", e).into()))?;

    let cursor = Cursor::new(binary_data);
    CsvReader::new(cursor)
        .finish()
}