feat: db query

app.py

@@ -2,19 +2,24 @@ from huggingface_hub import hf_hub_download
 from gensim.models import Word2Vec
 
 import faiss
+import duckdb
 
 import streamlit as st
 import pandas as pd
 import dask.dataframe as dd
 
 @st.cache_data
-def get_dask_df(df_path='bin/data.parquet'):
-    return dd.read_parquet(df_path)
+def get_db(path='arxiv.db'):
+    con = duckdb.connect(path)
 
 
-@st.cache_data
-def query_rows(rows: list):
-    df = get_dask_df()
+def query_neighbours(rows: list):
+    con = get_db()
+    placeholders = ",".join("?" for _ in rows)
+    return con.execute(
+        f"SELECT * FROM arxiv WHERE column0 IN ({placeholders})",
+        rows,
+    ).fetchall()
 
 @st.cache_data
 def get_model():