changes for ncruces go

This commit is contained in:
Alex Garcia 2024-07-11 22:36:18 -07:00
parent 374bfe3892
commit 73b9156a7c
7 changed files with 200 additions and 13 deletions

View file

@ -0,0 +1,35 @@
diff --git a/embed/build.sh b/embed/build.sh
index 36183bb..d202a26 100755
--- a/embed/build.sh
+++ b/embed/build.sh
@@ -20,6 +20,7 @@ WASI_SDK="$ROOT/tools/wasi-sdk-22.0/bin"
-Wl,--stack-first \
-Wl,--import-undefined \
-D_HAVE_SQLITE_CONFIG_H \
+ -DSQLITE_VEC_OMIT_FS=1 \
-DSQLITE_CUSTOM_INCLUDE=sqlite_opt.h \
$(awk '{print "-Wl,--export="$0}' exports.txt)
diff --git a/embed/sqlite3.wasm b/embed/sqlite3.wasm
index 569f0b3..b1d9693 100755
Binary files a/embed/sqlite3.wasm and b/embed/sqlite3.wasm differ
diff --git a/sqlite3/main.c b/sqlite3/main.c
index 4fb926b..49a9ee9 100644
--- a/sqlite3/main.c
+++ b/sqlite3/main.c
@@ -16,7 +16,7 @@
#include "time.c"
#include "vfs.c"
#include "vtab.c"
-
+#include "../../../../../sqlite-vec.c"
sqlite3_destructor_type malloc_destructor = &free;
__attribute__((constructor)) void init() {
@@ -28,4 +28,5 @@ __attribute__((constructor)) void init() {
sqlite3_auto_extension((void (*)(void))sqlite3_series_init);
sqlite3_auto_extension((void (*)(void))sqlite3_uint_init);
sqlite3_auto_extension((void (*)(void))sqlite3_time_init);
+ sqlite3_auto_extension((void (*)(void))sqlite3_vec_init);
}
\ No newline at end of file

2
examples/simple-go-ncruces/.gitignore vendored Normal file
View file

@ -0,0 +1,2 @@
demo
*.wasm

View file

@ -0,0 +1,6 @@
demo: demo.go
go build -o $@ $<
sqlite3.vec.wasm:
exit 1

View file

@ -0,0 +1,114 @@
package main
import (
"bytes"
_ "embed"
"encoding/binary"
"fmt"
"log"
"github.com/ncruces/go-sqlite3"
)
func serializeFloat32(vector []float32) ([]byte, error) {
buf := new(bytes.Buffer)
err := binary.Write(buf, binary.LittleEndian, vector)
if err != nil {
return nil, err
}
return buf.Bytes(), nil
}
const memory = ":memory:"
//go:embed sqlite3.vec.wasm
var sqliteWithVecWasm []byte
func main() {
sqlite3.Binary = sqliteWithVecWasm
db, err := sqlite3.Open(memory)
if err != nil {
log.Fatal(err)
}
stmt, _, err := db.Prepare(`SELECT sqlite_version(), vec_version()`)
if err != nil {
log.Fatal(err)
}
err = db.Exec("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[4])")
if err != nil {
log.Fatal(err)
}
items := map[int][]float32{
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},
}
q := []float32{0.3, 0.3, 0.3, 0.3}
stmt, _, err = db.Prepare("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")
if err != nil {
log.Fatal(err)
}
for id, values := range items {
v, err := serializeFloat32(values)
if err != nil {
log.Fatal(err)
}
stmt.BindInt(1, id)
stmt.BindBlob(2, v)
err = stmt.Exec()
if err != nil {
log.Fatal(err)
}
stmt.Reset()
}
stmt.Close()
stmt, _, err = db.Prepare(`
SELECT
rowid,
distance
FROM vec_items
WHERE embedding MATCH ?
ORDER BY distance
LIMIT 3
`);
if err != nil {
log.Fatal(err)
}
query, err := serializeFloat32(q)
if err != nil {
log.Fatal(err)
}
stmt.BindBlob(1, query)
for stmt.Step() {
rowid := stmt.ColumnInt64(0)
distance := stmt.ColumnFloat(1)
fmt.Printf("rowid=%d, distance=%f\n", rowid, distance)
}
if err := stmt.Err(); err != nil {
log.Fatal(err)
}
err = stmt.Close()
if err != nil {
log.Fatal(err)
}
err = db.Close()
if err != nil {
log.Fatal(err)
}
}

View file

@ -0,0 +1,13 @@
module asg017.com/ex1
go 1.22.5
require github.com/ncruces/go-sqlite3 v0.17.1
require (
github.com/ncruces/julianday v1.0.0 // indirect
github.com/tetratelabs/wazero v1.7.3 // indirect
golang.org/x/sys v0.22.0 // indirect
)
//replace github.com/ncruces/go-sqlite3 => ../go-sqlite3

View file

@ -0,0 +1,10 @@
github.com/ncruces/go-sqlite3 v0.17.1 h1:VxTjDpCn87FaFlKMaAYC1jP7ND0d4UNj+6G4IQDHbgI=
github.com/ncruces/go-sqlite3 v0.17.1/go.mod h1:FnCyui8SlDoL0mQZ5dTouNo7s7jXS0kJv9lBt1GlM9w=
github.com/ncruces/julianday v1.0.0 h1:fH0OKwa7NWvniGQtxdJRxAgkBMolni2BjDHaWTxqt7M=
github.com/ncruces/julianday v1.0.0/go.mod h1:Dusn2KvZrrovOMJuOt0TNXL6tB7U2E8kvza5fFc9G7g=
github.com/tetratelabs/wazero v1.7.3 h1:PBH5KVahrt3S2AHgEjKu4u+LlDbbk+nsGE3KLucy6Rw=
github.com/tetratelabs/wazero v1.7.3/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=
golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=

View file

@ -56,12 +56,14 @@ SQLITE_EXTENSION_INIT1
#ifndef _WIN32
#ifndef __EMSCRIPTEN__
#ifndef __COSMOPOLITAN__
#ifndef __wasi__
typedef u_int8_t uint8_t;
typedef u_int16_t uint16_t;
typedef u_int64_t uint64_t;
#endif
#endif
#endif
#endif
typedef int8_t i8;
typedef uint8_t u8;
@ -375,7 +377,7 @@ static f32 distance_hamming(const void *a, const void *b, const void *d) {
// from SQLite source:
// https://github.com/sqlite/sqlite/blob/a509a90958ddb234d1785ed7801880ccb18b497e/src/json.c#L153
static const char jsonIsSpace[] = {
static const char jsonIsSpaceX[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@ -390,7 +392,7 @@ static const char jsonIsSpace[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
};
#define jsonIsspace(x) (jsonIsSpace[(unsigned char)x])
#define jsonIsspace(x) (jsonIsSpaceX[(unsigned char)x])
typedef void (*vector_cleanup)(void *p);
@ -2411,6 +2413,7 @@ struct vec_npy_each_cursor {
static unsigned char NPY_MAGIC[6] = "\x93NUMPY";
#ifndef SQLITE_VEC_OMIT_FS
int parse_npy_file(sqlite3_vtab *pVTab, FILE *file, vec_npy_each_cursor *pCur) {
int n;
fseek(file, 0, SEEK_END);
@ -2499,6 +2502,7 @@ int parse_npy_file(sqlite3_vtab *pVTab, FILE *file, vec_npy_each_cursor *pCur) {
pCur->file = file;
return SQLITE_OK;
}
#endif
int parse_npy_buffer(sqlite3_vtab *pVTab, const unsigned char *buffer,
int bufferLength, void **data, size_t *numElements,
@ -2595,7 +2599,9 @@ static int vec_npy_eachOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) {
static int vec_npy_eachClose(sqlite3_vtab_cursor *cur) {
vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)cur;
if (pCur->file) {
#ifndef SQLITE_VEC_OMIT_FS
fclose(pCur->file);
#endif
pCur->file = NULL;
}
if (pCur->chunksBuffer) {
@ -2649,7 +2655,9 @@ static int vec_npy_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
vec_npy_each_cursor *pCur = (vec_npy_each_cursor *)pVtabCursor;
if (pCur->file) {
#ifndef SQLITE_VEC_OMIT_FS
fclose(pCur->file);
#endif
pCur->file = NULL;
}
if (pCur->chunksBuffer) {
@ -2662,6 +2670,7 @@ static int vec_npy_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
struct VecNpyFile *f = NULL;
#ifndef SQLITE_VEC_OMIT_FS
if ((f = sqlite3_value_pointer(argv[0], SQLITE_VEC_NPY_FILE_NAME))) {
FILE *file = fopen(f->path, "r");
if (!file) {
@ -2671,11 +2680,15 @@ static int vec_npy_eachFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
rc = parse_npy_file(pVtabCursor->pVtab, file, pCur);
if (rc != SQLITE_OK) {
#ifndef SQLITE_VEC_OMIT_FS
fclose(file);
#endif
return rc;
}
} else {
} else
#endif
{
const unsigned char *input = sqlite3_value_blob(argv[0]);
int inputLength = sqlite3_value_bytes(argv[0]);
@ -2722,6 +2735,7 @@ static int vec_npy_eachNext(sqlite3_vtab_cursor *cur) {
return SQLITE_OK;
}
#ifndef SQLITE_VEC_OMIT_FS
// else: input is a file
pCur->currentChunkIndex++;
if (pCur->currentChunkIndex >= pCur->currentChunkSize) {
@ -2734,6 +2748,7 @@ static int vec_npy_eachNext(sqlite3_vtab_cursor *cur) {
}
pCur->currentChunkIndex = 0;
}
#endif
return SQLITE_OK;
}
@ -4152,15 +4167,6 @@ void bitmap_clear(u8 *bitmap, i32 n) {
memset(bitmap, 0, n / CHAR_BIT);
}
void bitmap_debug(u8 *bitmap, i32 n) {
for (int i = 0; i < n; i++) {
printf("%d", bitmap_get(bitmap, i));
if (i > 0 && (i % 8 == 0))
printf("|");
}
printf("\n");
}
/**
* @brief Finds the minimum k items in distances, and writes the indicies to
* out.
@ -6700,10 +6706,10 @@ __declspec(dllexport)
&vec_static_blob_entriesModule,
static_blob_data, NULL);
assert(rc == SQLITE_OK);
return SQLITE_OK;
}
#ifndef SQLITE_VEC_OMIT_FS
#ifdef _WIN32
__declspec(dllexport)
#endif
@ -6716,6 +6722,7 @@ __declspec(dllexport)
NULL, vec_npy_file, NULL, NULL, NULL);
return rc;
}
#endif
#ifdef SQLITE_VEC_ENABLE_TRACE_ENTRYPOINT