example updates

This commit is contained in:
Alex Garcia 2024-05-11 00:21:44 -07:00
parent 12c96fc6fe
commit 85fac4c5ef
10 changed files with 112 additions and 76 deletions

View file

@ -1,5 +1,4 @@
import { Database } from "bun:sqlite"; import { Database } from "bun:sqlite";
Database.setCustomSQLite("/usr/local/opt/sqlite3/lib/libsqlite3.dylib"); Database.setCustomSQLite("/usr/local/opt/sqlite3/lib/libsqlite3.dylib");
const db = new Database(":memory:"); const db = new Database(":memory:");
@ -14,6 +13,15 @@ const { sqlite_version, vec_version } = db
console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`); console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`);
const items = [
[1, [0.1, 0.1, 0.1, 0.1]],
[2, [0.2, 0.2, 0.2, 0.2]],
[3, [0.3, 0.3, 0.3, 0.3]],
[4, [0.4, 0.4, 0.4, 0.4]],
[5, [0.5, 0.5, 0.5, 0.5]],
];
const query = [0.3, 0.3, 0.3, 0.3];
db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])"); db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])");
const insertStmt = db.prepare( const insertStmt = db.prepare(
@ -22,16 +30,12 @@ const insertStmt = db.prepare(
const insertVectors = db.transaction((items) => { const insertVectors = db.transaction((items) => {
for (const [id, vector] of items) { for (const [id, vector] of items) {
insertStmt.run(BigInt(id), vector); insertStmt.run(BigInt(id), new Float32Array(vector));
} }
}); });
insertVectors([ insertVectors(items);
[1, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])],
[2, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])],
]);
const query = new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]);
const rows = db const rows = db
.prepare( .prepare(
` `
@ -41,7 +45,7 @@ const rows = db
FROM vec_items FROM vec_items
WHERE embedding MATCH ? WHERE embedding MATCH ?
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
` `
) )
.all(query); .all(query);

View file

@ -22,21 +22,25 @@ int main(int argc, char *argv[]) {
printf("sqlite_version=%s, vec_version=%s\n", sqlite3_column_text(stmt, 0), sqlite3_column_text(stmt, 1)); printf("sqlite_version=%s, vec_version=%s\n", sqlite3_column_text(stmt, 0), sqlite3_column_text(stmt, 1));
sqlite3_finalize(stmt); sqlite3_finalize(stmt);
rc = sqlite3_prepare_v2(db, "CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])", -1, &stmt, NULL); static const struct {
sqlite3_int64 id;
float vector[4];
} items[] = {
{1, {0.1, 0.1, 0.1, 0.1}},
{2, {0.2, 0.2, 0.2, 0.2}},
{3, {0.3, 0.3, 0.3, 0.3}},
{4, {0.4, 0.4, 0.4, 0.4}},
{5, {0.5, 0.5, 0.5, 0.5}},
};
float query[4] = {0.3, 0.3, 0.3, 0.3};
rc = sqlite3_prepare_v2(db, "CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])", -1, &stmt, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
rc = sqlite3_step(stmt); rc = sqlite3_step(stmt);
assert(rc == SQLITE_DONE); assert(rc == SQLITE_DONE);
sqlite3_finalize(stmt); sqlite3_finalize(stmt);
static const struct {
sqlite3_int64 id;
float vector[8];
} items[] = {
{1, {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}},
{2, {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}},
};
rc = sqlite3_exec(db, "BEGIN", NULL, NULL, NULL); rc = sqlite3_exec(db, "BEGIN", NULL, NULL, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
rc = sqlite3_prepare_v2(db, "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)", -1, &stmt, NULL); rc = sqlite3_prepare_v2(db, "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)", -1, &stmt, NULL);
@ -52,7 +56,6 @@ int main(int argc, char *argv[]) {
rc = sqlite3_exec(db, "COMMIT", NULL, NULL, NULL); rc = sqlite3_exec(db, "COMMIT", NULL, NULL, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
rc = sqlite3_prepare_v2(db, rc = sqlite3_prepare_v2(db,
"SELECT " "SELECT "
" rowid, " " rowid, "
@ -60,11 +63,10 @@ int main(int argc, char *argv[]) {
"FROM vec_items " "FROM vec_items "
"WHERE embedding MATCH ?1 " "WHERE embedding MATCH ?1 "
"ORDER BY distance " "ORDER BY distance "
"LIMIT 5 " "LIMIT 3 "
, -1, &stmt, NULL); , -1, &stmt, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
float query[8] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8};
sqlite3_bind_blob(stmt, 1, query, sizeof(query), SQLITE_STATIC); sqlite3_bind_blob(stmt, 1, query, sizeof(query), SQLITE_STATIC);
while(1) { while(1) {

View file

@ -1,9 +1,9 @@
import { Database } from "jsr:@db/sqlite@0.11"; import { Database } from "jsr:@db/sqlite@0.11";
//import { loadablePath } from "npm:sqlite-vec"; import * as sqliteVec from "npm:sqlite-vec@0.0.1-alpha.3";
const db = new Database(":memory:"); const db = new Database(":memory:");
db.enableLoadExtension = true; db.enableLoadExtension = true;
db.loadExtension("../../dist/vec0"); sqliteVec.load(db);
db.enableLoadExtension = false; db.enableLoadExtension = false;
const [sqlite_version, vec_version] = db const [sqlite_version, vec_version] = db
@ -11,24 +11,29 @@ const [sqlite_version, vec_version] = db
.value<[string]>()!; .value<[string]>()!;
console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`); console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`);
db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])"); const items = [
[1, [0.1, 0.1, 0.1, 0.1]],
[2, [0.2, 0.2, 0.2, 0.2]],
[3, [0.3, 0.3, 0.3, 0.3]],
[4, [0.4, 0.4, 0.4, 0.4]],
[5, [0.5, 0.5, 0.5, 0.5]],
];
const query = [0.3, 0.3, 0.3, 0.3];
db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])");
const insertStmt = db.prepare( const insertStmt = db.prepare(
"INSERT INTO vec_items(rowid, embedding) VALUES (?1, vec_f32(?2))" "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)"
); );
const insertVectors = db.transaction((items) => { const insertVectors = db.transaction((items) => {
for (const [id, vector] of items) { for (const [id, vector] of items) {
insertStmt.run(BigInt(id), new Uint8Array(vector.buffer)); insertStmt.run(BigInt(id), new Uint8Array(new Float32Array(vector).buffer));
} }
}); });
insertVectors([ insertVectors(items);
[1, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])],
[2, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])],
]);
const query = new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]);
const rows = db const rows = db
.prepare( .prepare(
` `
@ -41,7 +46,7 @@ const rows = db
LIMIT 5 LIMIT 5
` `
) )
.all([new Uint8Array(query.buffer)]); .all([new Uint8Array(new Float32Array(query).buffer)]);
console.log(rows); console.log(rows);

View file

@ -38,15 +38,19 @@ func main() {
} }
fmt.Printf("sqlite_version=%s, vec_version=%s\n", sqliteVersion, vecVersion) fmt.Printf("sqlite_version=%s, vec_version=%s\n", sqliteVersion, vecVersion)
_, err = db.Exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])") _, err = db.Exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
items := map[int][]float32{ items := map[int][]float32{
1: {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, 1: {0.1, 0.1, 0.1, 0.1},
2: {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, 2: {0.2, 0.2, 0.2, 0.2},
3: {0.3, 0.3, 0.3, 0.3},
4: {0.4, 0.4, 0.4, 0.4},
5: {0.5, 0.5, 0.5, 0.5},
} }
q := []float32{0.3, 0.3, 0.3, 0.3}
for id, values := range items { for id, values := range items {
v, err := serializeFloat32(values) v, err := serializeFloat32(values)
@ -59,7 +63,6 @@ func main() {
} }
} }
q := []float32{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}
query, err := serializeFloat32(q) query, err := serializeFloat32(q)
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
@ -72,7 +75,7 @@ func main() {
FROM vec_items FROM vec_items
WHERE embedding MATCH ? WHERE embedding MATCH ?
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
`, query) `, query)
if err != nil { if err != nil {

View file

@ -12,24 +12,29 @@ const { sqlite_version, vec_version } = db
console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`); console.log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`);
db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])"); const items = [
[1, [0.1, 0.1, 0.1, 0.1]],
[2, [0.2, 0.2, 0.2, 0.2]],
[3, [0.3, 0.3, 0.3, 0.3]],
[4, [0.4, 0.4, 0.4, 0.4]],
[5, [0.5, 0.5, 0.5, 0.5]],
];
const query = [0.3, 0.3, 0.3, 0.3];
db.exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])");
const insertStmt = db.prepare( const insertStmt = db.prepare(
"INSERT INTO vec_items(rowid, embedding) VALUES (?, vec_f32(?))" "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)"
); );
const insertVectors = db.transaction((items) => { const insertVectors = db.transaction((items) => {
for (const [id, vector] of items) { for (const [id, vector] of items) {
insertStmt.run(BigInt(id), vector); insertStmt.run(BigInt(id), new Float32Array(vector));
} }
}); });
insertVectors([ insertVectors(items);
[1, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])],
[2, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])],
]);
const query = new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]);
const rows = db const rows = db
.prepare( .prepare(
` `
@ -39,9 +44,9 @@ const rows = db
FROM vec_items FROM vec_items
WHERE embedding MATCH ? WHERE embedding MATCH ?
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
` `
) )
.all(query); .all(new Float32Array(query));
console.log(rows); console.log(rows);

View file

@ -27,11 +27,14 @@ sqlite_version, vec_version = db.execute(
print(f"sqlite_version={sqlite_version}, vec_version={vec_version}") print(f"sqlite_version={sqlite_version}, vec_version={vec_version}")
db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])") db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
items = [ items = [
(1, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]), (1, [0.1, 0.1, 0.1, 0.1]),
(2, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]), (2, [0.2, 0.2, 0.2, 0.2]),
(3, [0.3, 0.3, 0.3, 0.3]),
(4, [0.4, 0.4, 0.4, 0.4]),
(5, [0.5, 0.5, 0.5, 0.5]),
] ]
with db: with db:
for item in items: for item in items:
@ -40,7 +43,7 @@ with db:
[item[0], serialize_f32(item[1])], [item[0], serialize_f32(item[1])],
) )
query = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] query = [0.3, 0.3, 0.3, 0.3]
rows = db.execute( rows = db.execute(
""" """
SELECT SELECT
@ -49,7 +52,7 @@ rows = db.execute(
FROM vec_items FROM vec_items
WHERE embedding MATCH ? WHERE embedding MATCH ?
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
""", """,
[serialize_f32(query)], [serialize_f32(query)],
).fetchall() ).fetchall()

View file

@ -10,19 +10,23 @@ db.enable_load_extension(false)
sqlite_version, vec_version = db.execute("select sqlite_version(), vec_version()").first sqlite_version, vec_version = db.execute("select sqlite_version(), vec_version()").first
puts "sqlite_version=#{sqlite_version}, vec_version=#{vec_version}" puts "sqlite_version=#{sqlite_version}, vec_version=#{vec_version}"
db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])") db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
items = [ items = [
[1, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]], [1, [0.1, 0.1, 0.1, 0.1]],
[2, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]] [2, [0.2, 0.2, 0.2, 0.2]],
[3, [0.3, 0.3, 0.3, 0.3]],
[4, [0.4, 0.4, 0.4, 0.4]],
[5, [0.5, 0.5, 0.5, 0.5]],
] ]
db.transaction do db.transaction do
items.each do |item| items.each do |item|
db.execute("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)", [item[0], item[1].pack("f*")]) db.execute("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)", [item[0], item[1].pack("f*")])
end end
end end
query = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] query = [0.3, 0.3, 0.3, 0.3]
rows = db.execute(<<-SQL, [query.pack("f*")]) rows = db.execute(<<-SQL, [query.pack("f*")])
SELECT SELECT
rowid, rowid,
@ -30,7 +34,7 @@ rows = db.execute(<<-SQL, [query.pack("f*")])
FROM vec_items FROM vec_items
WHERE embedding MATCH ? WHERE embedding MATCH ?
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
SQL SQL
puts rows puts rows

View file

@ -19,13 +19,16 @@ fn main() -> Result<()> {
println!("sqlite_version={sqlite_version}, vec_version={vec_version}"); println!("sqlite_version={sqlite_version}, vec_version={vec_version}");
let items: Vec<(usize, Vec<f32>)> = vec![ let items: Vec<(usize, Vec<f32>)> = vec![
(1, vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]), (1, vec![0.1, 0.1, 0.1, 0.1]),
(2, vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]), (2, vec![0.2, 0.2, 0.2, 0.2]),
(3, vec![0.3, 0.3, 0.3, 0.3]),
(4, vec![0.4, 0.4, 0.4, 0.4]),
(5, vec![0.5, 0.5, 0.5, 0.5]),
]; ];
println!("{x}"); println!("{x}");
db.execute( db.execute(
"CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8])", "CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])",
[], [],
)?; )?;
let mut stmt = db.prepare("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")?; let mut stmt = db.prepare("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")?;
@ -33,7 +36,7 @@ fn main() -> Result<()> {
stmt.execute(rusqlite::params![item.0, item.1.as_bytes()])?; stmt.execute(rusqlite::params![item.0, item.1.as_bytes()])?;
} }
let query: Vec<f32> = vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]; let query: Vec<f32> = vec![0.3, 0.3, 0.3, 0.3];
let result: Vec<(i64, f64)> = db let result: Vec<(i64, f64)> = db
.prepare( .prepare(
r" r"
@ -43,7 +46,7 @@ fn main() -> Result<()> {
FROM vec_items FROM vec_items
WHERE embedding MATCH ?1 WHERE embedding MATCH ?1
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
", ",
)? )?
.query_map([query.as_bytes()], |r| Ok((r.get(0)?, r.get(1)?)))? .query_map([query.as_bytes()], |r| Ok((r.get(0)?, r.get(1)?)))?

View file

@ -4,21 +4,24 @@
select sqlite_version(), vec_version(); select sqlite_version(), vec_version();
CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8]); CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4]);
INSERT INTO vec_items(rowid, embedding) INSERT INTO vec_items(rowid, embedding)
select select
value ->> 0, value ->> 0,
value ->> 1 value ->> 1
from json_each('[ from json_each('[
[1, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]], [1, [0.1, 0.1, 0.1, 0.1]],
[2, [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]] [2, [0.2, 0.2, 0.2, 0.2]],
[3, [0.3, 0.3, 0.3, 0.3]],
[4, [0.4, 0.4, 0.4, 0.4]],
[5, [0.5, 0.5, 0.5, 0.5]]
]'); ]');
SELECT SELECT
rowid, rowid,
distance distance
FROM vec_items FROM vec_items
WHERE embedding MATCH '[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]' WHERE embedding MATCH '[0.3, 0.3, 0.3, 0.3]'
ORDER BY distance ORDER BY distance
LIMIT 5; LIMIT 3;

View file

@ -5,33 +5,37 @@
<div id="target"></div> <div id="target"></div>
<script type="module"> <script type="module">
import {default as init} from "https://cdn.jsdelivr.net/npm/sqlite-vec-wasm-demo@latest/sqlite3.mjs"; import {default as init} from "https://cdn.jsdelivr.net/npm/sqlite-vec-wasm-demo@latest/sqlite3.mjs";
function log(msg) { function log(msg) {
const pre = document.querySelector('#target').appendChild(document.createElement("pre")); const pre = document.querySelector('#target').appendChild(document.createElement("pre"));
pre.textContent = msg; pre.textContent = msg;
} }
const sqlite3 = await init(); const sqlite3 = await init();
const db = new sqlite3.oo1.DB(":memory:"); const db = new sqlite3.oo1.DB(":memory:");
window.db = db;
const [sqlite_version, vec_version] = db.selectArray('select sqlite_version(), vec_version();') const [sqlite_version, vec_version] = db.selectArray('select sqlite_version(), vec_version();')
log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`); log(`sqlite_version=${sqlite_version}, vec_version=${vec_version}`);
const items = [ const items = [
[1, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])], [1, [0.1, 0.1, 0.1, 0.1]],
[2, new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])] [2, [0.2, 0.2, 0.2, 0.2]],
[3, [0.3, 0.3, 0.3, 0.3]],
[4, [0.4, 0.4, 0.4, 0.4]],
[5, [0.5, 0.5, 0.5, 0.5]],
]; ];
const query = [0.3, 0.3, 0.3, 0.3];
db.exec('CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[8]);'); db.exec('CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4]);');
const stmt = db.prepare("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)"); const stmt = db.prepare("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)");
for(const item of items) { for(const item of items) {
stmt stmt
.bind(1, item[0]) .bind(1, item[0])
.bind(2, item[1].buffer) .bind(2, new Float32Array(item[1]).buffer)
.stepReset(); .stepReset();
} }
stmt.finalize(); stmt.finalize();
const query = new Float32Array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]);
const rows = db const rows = db
.selectArrays( .selectArrays(
` `
@ -41,8 +45,8 @@
FROM vec_items FROM vec_items
WHERE embedding MATCH ? WHERE embedding MATCH ?
ORDER BY distance ORDER BY distance
LIMIT 5 LIMIT 3
`, query.buffer `, new Float32Array(query).buffer
); );
log(JSON.stringify(rows)); log(JSON.stringify(rows));