DuckDb python , example

add.py

import duckdb
import uuid

# ファイルに保存するデータベース
con = duckdb.connect(database='vector.db')

# INSTALL vss; LOAD vss;
con.execute("INSTALL vss; LOAD vss;")

create_sql = """
 CREATE TABLE IF NOT EXISTS embeddings (
      id TEXT PRIMARY KEY,
      text VARCHAR,
      vector FLOAT[3]
);
"""
print(create_sql)

add_vector = [0.1, 0.2, 0.3]
newUID = str(uuid.uuid4())

con.execute(create_sql)

# 
insert_sql = """
INSERT INTO embeddings 
      (id, text, vector) VALUES 
      (?, ?, CAST(? AS FLOAT[3]))
"""
print(insert_sql)

# 1件 登録
con.execute(insert_sql , (newUID , "test-text-1", add_vector))

search.py

import duckdb
import uuid

# ファイルに保存するデータベース
con = duckdb.connect(database='vector.db')

# INSTALL vss; LOAD vss;
con.execute("INSTALL vss; LOAD vss;")

query_vector = [0.1, 0.2, 0.31]
query_vec_str = [str(n) for n in query_vector]
#print(query_vec_str)

select_sql = """SELECT 
id, 
text, 
vector, 
array_cosine_similarity(vector, CAST(? AS FLOAT[3])) AS similarity
FROM embeddings
ORDER BY similarity DESC
LIMIT 2;
"""

resp = con.execute(select_sql, [query_vec_str]).fetchall()
#print(resp)
for row in resp:
    print(row[1])
    print("sim:"+ str(row[3]))