apunkt/sqlite-vec
1
0
Fork 0
mirror of https://github.com/asg017/sqlite-vec.git synced 2026-04-26 09:16:22 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

make format

Browse source
This commit is contained in:
Alex Garcia 2024-06-13 16:32:57 -07:00
parent dcb8bf5e53
commit df48ac2416
1 changed files with 280 additions and 267 deletions
Download patch file Download diff file Expand all files Collapse all files

371
sqlite-vec.c
View file

@ -1,6 +1,7 @@
#include "sqlite-vec.h" #include "sqlite-vec.h"
#include <assert.h> #include <assert.h>
#include <errno.h> #include <errno.h>
#include <float.h>
#include <inttypes.h> #include <inttypes.h>
#include <limits.h> #include <limits.h>
#include <math.h> #include <math.h>
@ -9,7 +10,6 @@
#include <stdio.h> #include <stdio.h>
#include <stdlib.h> #include <stdlib.h>
#include <string.h> #include <string.h>
#include <float.h>
#include "sqlite3ext.h" #include "sqlite3ext.h"
SQLITE_EXTENSION_INIT1 SQLITE_EXTENSION_INIT1
@ -186,7 +186,8 @@ static f32 l2_sqr_float_neon(const void *pVect1v, const void *pVect2v,
sum3 = vfmaq_f32(sum3, diff, diff); sum3 = vfmaq_f32(sum3, diff, diff);
} }
f32 sum_scalar = vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3))); f32 sum_scalar =
vaddvq_f32(vaddq_f32(vaddq_f32(sum0, sum1), vaddq_f32(sum2, sum3)));
const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4)); const f32 *pEnd2 = pVect1 + (qty - (qty16 << 4));
while (pVect1 < pEnd2) { while (pVect1 < pEnd2) {
f32 diff = *pVect1 - *pVect2; f32 diff = *pVect1 - *pVect2;
@ -284,11 +285,11 @@ static f32 distance_l2_sqr_float(const void *a, const void *b, const void *d) {
} }
static f32 distance_l2_sqr_int8(const void *a, const void *b, const void *d) { static f32 distance_l2_sqr_int8(const void *a, const void *b, const void *d) {
#ifdef SQLITE_VEC_ENABLE_NEON #ifdef SQLITE_VEC_ENABLE_NEON
if ((*(const size_t *)d) > 7) { if ((*(const size_t *)d) > 7) {
return l2_sqr_int8_neon(a, b, d); return l2_sqr_int8_neon(a, b, d);
} }
#endif #endif
return l2_sqr_int8(a, b, d); return l2_sqr_int8(a, b, d);
} }
@ -352,8 +353,8 @@ static f32 distance_hamming_u8(u8 *a, u8 *b, size_t n) {
} }
#ifdef _MSC_VER #ifdef _MSC_VER
# include <intrin.h> #include <intrin.h>
# define __builtin_popcountl __popcnt64 #define __builtin_popcountl __popcnt64
#endif #endif
static f32 distance_hamming_u64(u64 *a, u64 *b, size_t n) { static f32 distance_hamming_u64(u64 *a, u64 *b, size_t n) {
@ -419,7 +420,8 @@ struct Array {
* @param array * @param array
* @param element_size * @param element_size
* @param init_capacity * @param init_capacity
* @return SQLITE_OK on success, error code on failure. Only error is SQLITE_NOMEM * @return SQLITE_OK on success, error code on failure. Only error is
* SQLITE_NOMEM
*/ */
int array_init(struct Array *array, size_t element_size, size_t init_capacity) { int array_init(struct Array *array, size_t element_size, size_t init_capacity) {
void *z = sqlite3_malloc(element_size * init_capacity); void *z = sqlite3_malloc(element_size * init_capacity);
@ -493,7 +495,7 @@ static int fvec_from_value(sqlite3_value *value, f32 **vector,
struct Array x; struct Array x;
int rc = array_init(&x, sizeof(f32), ceil(source_len / 2.0)); int rc = array_init(&x, sizeof(f32), ceil(source_len / 2.0));
if(rc != SQLITE_OK) { if (rc != SQLITE_OK) {
return rc; return rc;
} }
@ -658,10 +660,8 @@ static int int8_vec_from_value(sqlite3_value *value, i8 **vector,
errno = 0; errno = 0;
long result = strtol(ptr, &endptr, 10); long result = strtol(ptr, &endptr, 10);
if ((errno != 0 && result == 0) if ((errno != 0 && result == 0) ||
|| (errno == ERANGE && (errno == ERANGE && (result == LONG_MAX || result == LONG_MIN))) {
(result == LONG_MAX || result == LONG_MIN))
) {
sqlite3_free(x.z); sqlite3_free(x.z);
*pzErr = sqlite3_mprintf("JSON parsing error"); *pzErr = sqlite3_mprintf("JSON parsing error");
return SQLITE_ERROR; return SQLITE_ERROR;
@ -678,7 +678,8 @@ static int int8_vec_from_value(sqlite3_value *value, i8 **vector,
if (result < INT8_MIN || result > INT8_MAX) { if (result < INT8_MIN || result > INT8_MAX) {
sqlite3_free(x.z); sqlite3_free(x.z);
*pzErr = sqlite3_mprintf("JSON parsing error: value out of range for int8"); *pzErr =
sqlite3_mprintf("JSON parsing error: value out of range for int8");
return SQLITE_ERROR; return SQLITE_ERROR;
} }
@ -1101,7 +1102,7 @@ static void vec_quantize_binary(sqlite3_context *context, int argc,
if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) { if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
u8 *out = sqlite3_malloc(dimensions / CHAR_BIT); u8 *out = sqlite3_malloc(dimensions / CHAR_BIT);
if(!out) { if (!out) {
cleanup(vector); cleanup(vector);
sqlite3_result_error_code(context, SQLITE_NOMEM); sqlite3_result_error_code(context, SQLITE_NOMEM);
return; return;
@ -1114,7 +1115,7 @@ static void vec_quantize_binary(sqlite3_context *context, int argc,
sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT); sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
} else if (elementType == SQLITE_VEC_ELEMENT_TYPE_INT8) { } else if (elementType == SQLITE_VEC_ELEMENT_TYPE_INT8) {
u8 *out = sqlite3_malloc(dimensions / CHAR_BIT); u8 *out = sqlite3_malloc(dimensions / CHAR_BIT);
if(!out) { if (!out) {
cleanup(vector); cleanup(vector);
sqlite3_result_error_code(context, SQLITE_NOMEM); sqlite3_result_error_code(context, SQLITE_NOMEM);
return; return;
@ -1126,7 +1127,8 @@ static void vec_quantize_binary(sqlite3_context *context, int argc,
sqlite3_result_blob(context, out, dimensions / CHAR_BIT, sqlite3_free); sqlite3_result_blob(context, out, dimensions / CHAR_BIT, sqlite3_free);
sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT); sqlite3_result_subtype(context, SQLITE_VEC_ELEMENT_TYPE_BIT);
} else { } else {
sqlite3_result_error(context, "Can only binary quantize float or int8 vectors", -1); sqlite3_result_error(context,
"Can only binary quantize float or int8 vectors", -1);
return; return;
} }
} }
@ -1371,10 +1373,9 @@ static void vec_to_json(sqlite3_context *context, int argc,
} }
if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) { if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
f32 value = ((f32 *)vector)[i]; f32 value = ((f32 *)vector)[i];
if(isnan(value)) { if (isnan(value)) {
sqlite3_str_appendall(str, "null"); sqlite3_str_appendall(str, "null");
} } else {
else {
sqlite3_str_appendf(str, "%f", value); sqlite3_str_appendf(str, "%f", value);
} }
@ -1422,7 +1423,7 @@ static void vec_normalize(sqlite3_context *context, int argc,
} }
f32 *out = sqlite3_malloc(dimensions * sizeof(f32)); f32 *out = sqlite3_malloc(dimensions * sizeof(f32));
if(!out) { if (!out) {
cleanup(vector); cleanup(vector);
sqlite3_result_error_code(context, SQLITE_NOMEM); sqlite3_result_error_code(context, SQLITE_NOMEM);
return; return;
@ -1975,10 +1976,10 @@ static sqlite3_module vec_eachModule = {
/* xRelease */ 0, /* xRelease */ 0,
/* xRollbackTo */ 0, /* xRollbackTo */ 0,
/* xShadowName */ 0, /* xShadowName */ 0,
#if SQLITE_VERSION_NUMBER >= 3440000 #if SQLITE_VERSION_NUMBER >= 3440000
/* xIntegrity */ 0 /* xIntegrity */ 0
#endif #endif
}; };
#pragma endregion #pragma endregion
@ -2583,10 +2584,10 @@ static sqlite3_module vec_npy_eachModule = {
/* xRelease */ 0, /* xRelease */ 0,
/* xRollbackTo */ 0, /* xRollbackTo */ 0,
/* xShadowName */ 0, /* xShadowName */ 0,
#if SQLITE_VERSION_NUMBER >= 3440000 #if SQLITE_VERSION_NUMBER >= 3440000
/* xIntegrity */ 0, /* xIntegrity */ 0,
#endif #endif
}; };
#pragma endregion #pragma endregion
@ -2918,22 +2919,22 @@ int vec0_new_chunk(vec0_vtab *p, i64 *chunk_rowid) {
"(size, validity, rowids) " "(size, validity, rowids) "
"VALUES (?, ?, ?);", "VALUES (?, ?, ?);",
p->schemaName, p->tableName); p->schemaName, p->tableName);
if(!zSql) { if (!zSql) {
return SQLITE_NOMEM; return SQLITE_NOMEM;
} }
rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
sqlite3_free(zSql); sqlite3_free(zSql);
if(rc != SQLITE_OK) { if (rc != SQLITE_OK) {
return rc; return rc;
} }
#ifdef SQLITE_VEC_THREADSAFE #ifdef SQLITE_VEC_THREADSAFE
sqlite3_mutex_enter(sqlite3_db_mutex(p->db)); sqlite3_mutex_enter(sqlite3_db_mutex(p->db));
#endif #endif
rc = sqlite3_bind_int64(stmt, 1, p->chunk_size); // size rc = sqlite3_bind_int64(stmt, 1, p->chunk_size); // size
if(rc != SQLITE_OK) { if (rc != SQLITE_OK) {
#ifdef SQLITE_VEC_THREADSAFE #ifdef SQLITE_VEC_THREADSAFE
sqlite3_mutex_leave(sqlite3_db_mutex(p->db)); sqlite3_mutex_leave(sqlite3_db_mutex(p->db));
#endif #endif
sqlite3_finalize(stmt); sqlite3_finalize(stmt);
return SQLITE_ERROR; return SQLITE_ERROR;
} }
@ -4002,8 +4003,7 @@ static int vec0Eof(sqlite3_vtab_cursor *cur) {
case SQLITE_VEC0_QUERYPLAN_KNN: { case SQLITE_VEC0_QUERYPLAN_KNN: {
todo_assert(pCur->knn_data); todo_assert(pCur->knn_data);
return (pCur->knn_data->current_idx >= pCur->knn_data->k) || return (pCur->knn_data->current_idx >= pCur->knn_data->k) ||
(pCur->knn_data->distances[pCur->knn_data->current_idx] == (pCur->knn_data->distances[pCur->knn_data->current_idx] == FLT_MAX);
FLT_MAX);
} }
case SQLITE_VEC0_QUERYPLAN_POINT: { case SQLITE_VEC0_QUERYPLAN_POINT: {
todo_assert(pCur->point_data); todo_assert(pCur->point_data);
@ -4641,29 +4641,31 @@ static sqlite3_module vec0Module = {
/* xRelease */ 0, /* xRelease */ 0,
/* xRollbackTo */ 0, /* xRollbackTo */ 0,
/* xShadowName */ vec0ShadowName, /* xShadowName */ vec0ShadowName,
#if SQLITE_VERSION_NUMBER >= 3440000 #if SQLITE_VERSION_NUMBER >= 3440000
/* xIntegrity */ 0, // TODO /* xIntegrity */ 0, // TODO
#endif #endif
}; };
#pragma endregion #pragma endregion
#ifdef SQLITE_VEC_ENABLE_EXPERIMENTAL #ifdef SQLITE_VEC_ENABLE_EXPERIMENTAL
static char * POINTER_NAME_STATIC_BLOB_DEF = "vec0-static_blob_def"; static char *POINTER_NAME_STATIC_BLOB_DEF = "vec0-static_blob_def";
struct static_blob_definition { struct static_blob_definition {
void * p; void *p;
size_t dimensions; size_t dimensions;
size_t nvectors; size_t nvectors;
enum VectorElementType element_type; enum VectorElementType element_type;
}; };
static void vec_static_blob_from_raw(sqlite3_context *context, int argc, sqlite3_value **argv) { static void vec_static_blob_from_raw(sqlite3_context *context, int argc,
struct static_blob_definition * p; sqlite3_value **argv) {
struct static_blob_definition *p;
p = sqlite3_malloc(sizeof(*p)); p = sqlite3_malloc(sizeof(*p));
todo_assert(p); todo_assert(p);
p->p = sqlite3_value_int64(argv[0]); p->p = sqlite3_value_int64(argv[0]);
p->element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32; p->element_type = SQLITE_VEC_ELEMENT_TYPE_FLOAT32;
p->dimensions = sqlite3_value_int64(argv[2]); p->dimensions = sqlite3_value_int64(argv[2]);
p->nvectors = sqlite3_value_int64(argv[3]); p->nvectors = sqlite3_value_int64(argv[3]);
sqlite3_result_pointer(context, p, POINTER_NAME_STATIC_BLOB_DEF, sqlite3_free); sqlite3_result_pointer(context, p, POINTER_NAME_STATIC_BLOB_DEF,
sqlite3_free);
} }
#pragma region vec_static_blobs() table function #pragma region vec_static_blobs() table function
@ -4671,8 +4673,8 @@ static void vec_static_blob_from_raw(sqlite3_context *context, int argc, sqlite3
typedef struct static_blob static_blob; typedef struct static_blob static_blob;
struct static_blob { struct static_blob {
char * name; char *name;
void * p; void *p;
size_t dimensions; size_t dimensions;
size_t nvectors; size_t nvectors;
enum VectorElementType element_type; enum VectorElementType element_type;
@ -4686,7 +4688,7 @@ struct vec_static_blob_data {
typedef struct vec_static_blobs_vtab vec_static_blobs_vtab; typedef struct vec_static_blobs_vtab vec_static_blobs_vtab;
struct vec_static_blobs_vtab { struct vec_static_blobs_vtab {
sqlite3_vtab base; sqlite3_vtab base;
vec_static_blob_data * data; vec_static_blob_data *data;
}; };
typedef struct vec_static_blobs_cursor vec_static_blobs_cursor; typedef struct vec_static_blobs_cursor vec_static_blobs_cursor;
@ -4696,8 +4698,8 @@ struct vec_static_blobs_cursor {
}; };
static int vec_static_blobsConnect(sqlite3 *db, void *pAux, int argc, static int vec_static_blobsConnect(sqlite3 *db, void *pAux, int argc,
const char *const *argv, sqlite3_vtab **ppVtab, const char *const *argv,
char **pzErr) { sqlite3_vtab **ppVtab, char **pzErr) {
vec_static_blobs_vtab *pNew; vec_static_blobs_vtab *pNew;
#define VEC_STATIC_BLOBS_NAME 0 #define VEC_STATIC_BLOBS_NAME 0
#define VEC_STATIC_BLOBS_DATA 1 #define VEC_STATIC_BLOBS_DATA 1
@ -4722,26 +4724,28 @@ static int vec_static_blobsDisconnect(sqlite3_vtab *pVtab) {
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blobsUpdate(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv, static int vec_static_blobsUpdate(sqlite3_vtab *pVTab, int argc,
sqlite_int64 *pRowid) { sqlite3_value **argv, sqlite_int64 *pRowid) {
vec_static_blobs_vtab* p = (vec_static_blobs_vtab*) pVTab; vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pVTab;
// DELETE operation // DELETE operation
if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) { if (argc == 1 && sqlite3_value_type(argv[0]) != SQLITE_NULL) {
return SQLITE_ERROR; return SQLITE_ERROR;
} }
// INSERT operation // INSERT operation
else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) { else if (argc > 1 && sqlite3_value_type(argv[0]) == SQLITE_NULL) {
const char * key = sqlite3_value_text(argv[2 + VEC_STATIC_BLOBS_NAME]); const char *key = sqlite3_value_text(argv[2 + VEC_STATIC_BLOBS_NAME]);
int idx = -1; int idx = -1;
for(int i = 0; i < MAX_STATIC_BLOBS; i++) { for (int i = 0; i < MAX_STATIC_BLOBS; i++) {
if(!p->data->static_blobs[i].name) { if (!p->data->static_blobs[i].name) {
p->data->static_blobs[i].name = sqlite3_mprintf("%s", key); p->data->static_blobs[i].name = sqlite3_mprintf("%s", key);
idx = i; idx = i;
break; break;
} }
} }
if(idx < 0) abort(); if (idx < 0)
struct static_blob_definition * def = sqlite3_value_pointer(argv[2 + VEC_STATIC_BLOBS_DATA], POINTER_NAME_STATIC_BLOB_DEF); abort();
struct static_blob_definition *def = sqlite3_value_pointer(
argv[2 + VEC_STATIC_BLOBS_DATA], POINTER_NAME_STATIC_BLOB_DEF);
p->data->static_blobs[idx].p = def->p; p->data->static_blobs[idx].p = def->p;
p->data->static_blobs[idx].dimensions = def->dimensions; p->data->static_blobs[idx].dimensions = def->dimensions;
p->data->static_blobs[idx].nvectors = def->nvectors; p->data->static_blobs[idx].nvectors = def->nvectors;
@ -4756,7 +4760,8 @@ static int vec_static_blobsUpdate(sqlite3_vtab *pVTab, int argc, sqlite3_value *
return SQLITE_ERROR; return SQLITE_ERROR;
} }
static int vec_static_blobsOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) { static int vec_static_blobsOpen(sqlite3_vtab *p,
sqlite3_vtab_cursor **ppCursor) {
vec_static_blobs_cursor *pCur; vec_static_blobs_cursor *pCur;
pCur = sqlite3_malloc(sizeof(*pCur)); pCur = sqlite3_malloc(sizeof(*pCur));
if (pCur == 0) if (pCur == 0)
@ -4772,7 +4777,8 @@ static int vec_static_blobsClose(sqlite3_vtab_cursor *cur) {
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blobsBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { static int vec_static_blobsBestIndex(sqlite3_vtab *pVTab,
sqlite3_index_info *pIdxInfo) {
pIdxInfo->idxNum = 1; pIdxInfo->idxNum = 1;
pIdxInfo->estimatedCost = (double)10; pIdxInfo->estimatedCost = (double)10;
pIdxInfo->estimatedRows = 10; pIdxInfo->estimatedRows = 10;
@ -4781,14 +4787,16 @@ static int vec_static_blobsBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pI
static int vec_static_blobsNext(sqlite3_vtab_cursor *cur); static int vec_static_blobsNext(sqlite3_vtab_cursor *cur);
static int vec_static_blobsFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, static int vec_static_blobsFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
const char *idxStr, int argc, sqlite3_value **argv) { const char *idxStr, int argc,
sqlite3_value **argv) {
vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)pVtabCursor; vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)pVtabCursor;
pCur->iRowid = -1; pCur->iRowid = -1;
vec_static_blobsNext(pVtabCursor); vec_static_blobsNext(pVtabCursor);
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blobsRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) { static int vec_static_blobsRowid(sqlite3_vtab_cursor *cur,
sqlite_int64 *pRowid) {
vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
*pRowid = pCur->iRowid; *pRowid = pCur->iRowid;
return SQLITE_OK; return SQLITE_OK;
@ -4796,10 +4804,10 @@ static int vec_static_blobsRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid)
static int vec_static_blobsNext(sqlite3_vtab_cursor *cur) { static int vec_static_blobsNext(sqlite3_vtab_cursor *cur) {
vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
vec_static_blobs_vtab * p = (vec_static_blobs_vtab *) pCur->base.pVtab; vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)pCur->base.pVtab;
pCur->iRowid++; pCur->iRowid++;
while(pCur->iRowid < MAX_STATIC_BLOBS) { while (pCur->iRowid < MAX_STATIC_BLOBS) {
if(p->data->static_blobs[pCur->iRowid].name) { if (p->data->static_blobs[pCur->iRowid].name) {
return SQLITE_OK; return SQLITE_OK;
} }
pCur->iRowid++; pCur->iRowid++;
@ -4812,19 +4820,21 @@ static int vec_static_blobsEof(sqlite3_vtab_cursor *cur) {
return pCur->iRowid >= MAX_STATIC_BLOBS; return pCur->iRowid >= MAX_STATIC_BLOBS;
} }
static int vec_static_blobsColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context, static int vec_static_blobsColumn(sqlite3_vtab_cursor *cur,
int i) { sqlite3_context *context, int i) {
vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur; vec_static_blobs_cursor *pCur = (vec_static_blobs_cursor *)cur;
vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)cur->pVtab; vec_static_blobs_vtab *p = (vec_static_blobs_vtab *)cur->pVtab;
switch (i) { switch (i) {
case VEC_STATIC_BLOBS_NAME: case VEC_STATIC_BLOBS_NAME:
sqlite3_result_text(context, p->data->static_blobs[pCur->iRowid].name, -1, SQLITE_TRANSIENT); sqlite3_result_text(context, p->data->static_blobs[pCur->iRowid].name, -1,
SQLITE_TRANSIENT);
break; break;
case VEC_STATIC_BLOBS_DATA: case VEC_STATIC_BLOBS_DATA:
sqlite3_result_null(context); sqlite3_result_null(context);
break; break;
case VEC_STATIC_BLOBS_DIMENSIONS: case VEC_STATIC_BLOBS_DIMENSIONS:
sqlite3_result_int64(context, p->data->static_blobs[pCur->iRowid].dimensions); sqlite3_result_int64(context,
p->data->static_blobs[pCur->iRowid].dimensions);
break; break;
case VEC_STATIC_BLOBS_COUNT: case VEC_STATIC_BLOBS_COUNT:
sqlite3_result_int64(context, p->data->static_blobs[pCur->iRowid].nvectors); sqlite3_result_int64(context, p->data->static_blobs[pCur->iRowid].nvectors);
@ -4833,7 +4843,6 @@ static int vec_static_blobsColumn(sqlite3_vtab_cursor *cur, sqlite3_context *con
return SQLITE_OK; return SQLITE_OK;
} }
static sqlite3_module vec_static_blobsModule = { static sqlite3_module vec_static_blobsModule = {
/* iVersion */ 3, /* iVersion */ 3,
/* xCreate */ 0, /* xCreate */ 0,
@ -4861,13 +4870,12 @@ static sqlite3_module vec_static_blobsModule = {
/* xShadowName */ 0}; /* xShadowName */ 0};
#pragma endregion #pragma endregion
#pragma region vec_static_blob_entries() table function #pragma region vec_static_blob_entries() table function
typedef struct vec_static_blob_entries_vtab vec_static_blob_entries_vtab; typedef struct vec_static_blob_entries_vtab vec_static_blob_entries_vtab;
struct vec_static_blob_entries_vtab { struct vec_static_blob_entries_vtab {
sqlite3_vtab base; sqlite3_vtab base;
static_blob * blob; static_blob *blob;
}; };
typedef enum { typedef enum {
VEC_SBE__QUERYPLAN_FULLSCAN = 1, VEC_SBE__QUERYPLAN_FULLSCAN = 1,
@ -4879,23 +4887,25 @@ struct vec_static_blob_entries_cursor {
sqlite3_vtab_cursor base; sqlite3_vtab_cursor base;
sqlite3_int64 iRowid; sqlite3_int64 iRowid;
vec_sbe_query_plan query_plan; vec_sbe_query_plan query_plan;
struct vec0_query_knn_data * knn_data; struct vec0_query_knn_data *knn_data;
}; };
static int vec_static_blob_entriesConnect(sqlite3 *db, void *pAux, int argc, static int vec_static_blob_entriesConnect(sqlite3 *db, void *pAux, int argc,
const char *const *argv, sqlite3_vtab **ppVtab, const char *const *argv,
char **pzErr) { sqlite3_vtab **ppVtab, char **pzErr) {
vec_static_blob_data * blob_data = pAux; vec_static_blob_data *blob_data = pAux;
int idx = -1; int idx = -1;
for(int i = 0; i < MAX_STATIC_BLOBS; i++) { for (int i = 0; i < MAX_STATIC_BLOBS; i++) {
if(!blob_data->static_blobs[i].name) continue; if (!blob_data->static_blobs[i].name)
if(strncmp(blob_data->static_blobs[i].name, argv[3], strlen(blob_data->static_blobs[i].name))==0) { continue;
if (strncmp(blob_data->static_blobs[i].name, argv[3],
strlen(blob_data->static_blobs[i].name)) == 0) {
idx = i; idx = i;
break; break;
} }
} }
if(idx < 0) abort(); if (idx < 0)
abort();
vec_static_blob_entries_vtab *pNew; vec_static_blob_entries_vtab *pNew;
#define VEC_STATIC_BLOB_ENTRIES_VECTOR 0 #define VEC_STATIC_BLOB_ENTRIES_VECTOR 0
#define VEC_STATIC_BLOB_ENTRIES_DISTANCE 1 #define VEC_STATIC_BLOB_ENTRIES_DISTANCE 1
@ -4914,10 +4924,10 @@ static int vec_static_blob_entriesConnect(sqlite3 *db, void *pAux, int argc,
} }
static int vec_static_blob_entriesCreate(sqlite3 *db, void *pAux, int argc, static int vec_static_blob_entriesCreate(sqlite3 *db, void *pAux, int argc,
const char *const *argv, sqlite3_vtab **ppVtab, const char *const *argv,
char **pzErr) { sqlite3_vtab **ppVtab, char **pzErr) {
vec_static_blob_entriesConnect(db, pAux, argc, argv, ppVtab, pzErr); vec_static_blob_entriesConnect(db, pAux, argc, argv, ppVtab, pzErr);
} }
static int vec_static_blob_entriesDisconnect(sqlite3_vtab *pVtab) { static int vec_static_blob_entriesDisconnect(sqlite3_vtab *pVtab) {
vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVtab; vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVtab;
@ -4925,7 +4935,8 @@ static int vec_static_blob_entriesDisconnect(sqlite3_vtab *pVtab) {
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blob_entriesOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor) { static int vec_static_blob_entriesOpen(sqlite3_vtab *p,
sqlite3_vtab_cursor **ppCursor) {
vec_static_blob_entries_cursor *pCur; vec_static_blob_entries_cursor *pCur;
pCur = sqlite3_malloc(sizeof(*pCur)); pCur = sqlite3_malloc(sizeof(*pCur));
if (pCur == 0) if (pCur == 0)
@ -4941,7 +4952,8 @@ static int vec_static_blob_entriesClose(sqlite3_vtab_cursor *cur) {
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab,
sqlite3_index_info *pIdxInfo) {
vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVTab; vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pVTab;
int iMatchTerm = -1; int iMatchTerm = -1;
int iLimitTerm = -1; int iLimitTerm = -1;
@ -4954,7 +4966,8 @@ static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, sqlite3_index_i
int iColumn = pIdxInfo->aConstraint[i].iColumn; int iColumn = pIdxInfo->aConstraint[i].iColumn;
int op = pIdxInfo->aConstraint[i].op; int op = pIdxInfo->aConstraint[i].op;
if (op == SQLITE_INDEX_CONSTRAINT_MATCH && iColumn == VEC_STATIC_BLOB_ENTRIES_VECTOR) { if (op == SQLITE_INDEX_CONSTRAINT_MATCH &&
iColumn == VEC_STATIC_BLOB_ENTRIES_VECTOR) {
if (iMatchTerm > -1) { if (iMatchTerm > -1) {
// TODO only 1 match operator at a time // TODO only 1 match operator at a time
return SQLITE_ERROR; return SQLITE_ERROR;
@ -4964,11 +4977,12 @@ static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, sqlite3_index_i
if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) { if (op == SQLITE_INDEX_CONSTRAINT_LIMIT) {
iLimitTerm = i; iLimitTerm = i;
} }
if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == VEC_STATIC_BLOB_ENTRIES_K) { if (op == SQLITE_INDEX_CONSTRAINT_EQ &&
iColumn == VEC_STATIC_BLOB_ENTRIES_K) {
iKTerm = i; iKTerm = i;
} }
} }
if(iMatchTerm >= 0) { if (iMatchTerm >= 0) {
if (iLimitTerm < 0 && iKTerm < 0) { if (iLimitTerm < 0 && iKTerm < 0) {
// TODO: error, match on vector1 should require a limit for KNN // TODO: error, match on vector1 should require a limit for KNN
return SQLITE_ERROR; return SQLITE_ERROR;
@ -4997,7 +5011,7 @@ static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, sqlite3_index_i
pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN; pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN;
pIdxInfo->estimatedCost = (double)10; // TODO vtab_value(?) as hint? pIdxInfo->estimatedCost = (double)10; // TODO vtab_value(?) as hint?
pIdxInfo->estimatedRows = 10;// TODO vtab_value(?) as hint? pIdxInfo->estimatedRows = 10; // TODO vtab_value(?) as hint?
pIdxInfo->orderByConsumed = 1; pIdxInfo->orderByConsumed = 1;
pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1; pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1;
@ -5010,8 +5024,7 @@ static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, sqlite3_index_i
pIdxInfo->aConstraintUsage[iKTerm].omit = 1; pIdxInfo->aConstraintUsage[iKTerm].omit = 1;
} }
} } else {
else {
pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN; pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN;
pIdxInfo->estimatedCost = (double)p->blob->nvectors; pIdxInfo->estimatedCost = (double)p->blob->nvectors;
pIdxInfo->estimatedRows = p->blob->nvectors; pIdxInfo->estimatedRows = p->blob->nvectors;
@ -5019,12 +5032,15 @@ static int vec_static_blob_entriesBestIndex(sqlite3_vtab *pVTab, sqlite3_index_i
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor,
const char *idxStr, int argc, sqlite3_value **argv) { int idxNum, const char *idxStr,
vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)pVtabCursor; int argc, sqlite3_value **argv) {
vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)pCur->base.pVtab; vec_static_blob_entries_cursor *pCur =
(vec_static_blob_entries_cursor *)pVtabCursor;
vec_static_blob_entries_vtab *p =
(vec_static_blob_entries_vtab *)pCur->base.pVtab;
if(idxNum == VEC_SBE__QUERYPLAN_KNN) { if (idxNum == VEC_SBE__QUERYPLAN_KNN) {
pCur->query_plan = VEC_SBE__QUERYPLAN_KNN; pCur->query_plan = VEC_SBE__QUERYPLAN_KNN;
struct vec0_query_knn_data *knn_data = struct vec0_query_knn_data *knn_data =
sqlite3_malloc(sizeof(struct vec0_query_knn_data)); sqlite3_malloc(sizeof(struct vec0_query_knn_data));
@ -5038,11 +5054,12 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, int i
enum VectorElementType elementType; enum VectorElementType elementType;
vector_cleanup cleanup; vector_cleanup cleanup;
char *err; char *err;
int rc = vector_from_value(argv[0], &queryVector, &dimensions, &elementType, &cleanup, &err); int rc = vector_from_value(argv[0], &queryVector, &dimensions, &elementType,
&cleanup, &err);
todo_assert(elementType == p->blob->element_type); todo_assert(elementType == p->blob->element_type);
todo_assert(dimensions == p->blob->dimensions); todo_assert(dimensions == p->blob->dimensions);
#define min(a, b) (((a) < (b)) ? (a) : (b)) #define min(a, b) (((a) < (b)) ? (a) : (b))
i64 k = min(sqlite3_value_int64(argv[1]), p->blob->nvectors); i64 k = min(sqlite3_value_int64(argv[1]), p->blob->nvectors);
todo_assert(k >= 0); todo_assert(k >= 0);
@ -5057,9 +5074,10 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, int i
f32 *distances = sqlite3_malloc(p->blob->nvectors * sizeof(f32)); f32 *distances = sqlite3_malloc(p->blob->nvectors * sizeof(f32));
todo_assert(distances); todo_assert(distances);
for(size_t i = 0; i < p->blob->nvectors; i++) { for (size_t i = 0; i < p->blob->nvectors; i++) {
float * v = ((float *) p->blob->p) + (i * p->blob->dimensions); float *v = ((float *)p->blob->p) + (i * p->blob->dimensions);
distances[i] = distance_l2_sqr_float(v, (float *) queryVector, &p->blob->dimensions); distances[i] =
distance_l2_sqr_float(v, (float *)queryVector, &p->blob->dimensions);
} }
min_idx(distances, k, topk_rowids, k); min_idx(distances, k, topk_rowids, k);
knn_data->current_idx = 0; knn_data->current_idx = 0;
@ -5068,8 +5086,7 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, int i
knn_data->rowids = topk_rowids; knn_data->rowids = topk_rowids;
pCur->knn_data = knn_data; pCur->knn_data = knn_data;
} } else {
else {
pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN; pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN;
pCur->iRowid = 0; pCur->iRowid = 0;
} }
@ -5077,7 +5094,8 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, int i
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_static_blob_entriesRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) { static int vec_static_blob_entriesRowid(sqlite3_vtab_cursor *cur,
sqlite_int64 *pRowid) {
vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
*pRowid = pCur->iRowid; *pRowid = pCur->iRowid;
return SQLITE_OK; return SQLITE_OK;
@ -5085,7 +5103,7 @@ static int vec_static_blob_entriesRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *
static int vec_static_blob_entriesNext(sqlite3_vtab_cursor *cur) { static int vec_static_blob_entriesNext(sqlite3_vtab_cursor *cur) {
vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
switch(pCur->query_plan) { switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
pCur->iRowid++; pCur->iRowid++;
return SQLITE_OK; return SQLITE_OK;
@ -5095,29 +5113,28 @@ static int vec_static_blob_entriesNext(sqlite3_vtab_cursor *cur) {
return SQLITE_OK; return SQLITE_OK;
} }
} }
} }
static int vec_static_blob_entriesEof(sqlite3_vtab_cursor *cur) { static int vec_static_blob_entriesEof(sqlite3_vtab_cursor *cur) {
vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
vec_static_blob_entries_vtab * p = (vec_static_blob_entries_vtab *) pCur->base.pVtab; vec_static_blob_entries_vtab *p =
switch(pCur->query_plan) { (vec_static_blob_entries_vtab *)pCur->base.pVtab;
switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
return (size_t) pCur->iRowid >= p->blob->nvectors; return (size_t)pCur->iRowid >= p->blob->nvectors;
} }
case VEC_SBE__QUERYPLAN_KNN: { case VEC_SBE__QUERYPLAN_KNN: {
return pCur->knn_data->current_idx >= pCur->knn_data->k; return pCur->knn_data->current_idx >= pCur->knn_data->k;
} }
} }
} }
static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context, static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur,
int i) { sqlite3_context *context, int i) {
vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur;
vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)cur->pVtab; vec_static_blob_entries_vtab *p = (vec_static_blob_entries_vtab *)cur->pVtab;
switch(pCur->query_plan) { switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
switch (i) { switch (i) {
case VEC_STATIC_BLOB_ENTRIES_VECTOR: case VEC_STATIC_BLOB_ENTRIES_VECTOR:
@ -5125,10 +5142,8 @@ static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur, sqlite3_conte
sqlite3_result_blob( sqlite3_result_blob(
context, context,
p->blob->p + (pCur->iRowid * p->blob->dimensions * sizeof(float)), p->blob->p + (pCur->iRowid * p->blob->dimensions * sizeof(float)),
p->blob->dimensions * sizeof(float), p->blob->dimensions * sizeof(float), SQLITE_STATIC);
SQLITE_STATIC sqlite3_result_subtype(context, p->blob->element_type);
);
sqlite3_result_subtype(context,p->blob->element_type);
break; break;
} }
return SQLITE_OK; return SQLITE_OK;
@ -5136,15 +5151,12 @@ static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur, sqlite3_conte
case VEC_SBE__QUERYPLAN_KNN: { case VEC_SBE__QUERYPLAN_KNN: {
switch (i) { switch (i) {
case VEC_STATIC_BLOB_ENTRIES_VECTOR: { case VEC_STATIC_BLOB_ENTRIES_VECTOR: {
i32 rowid = ((i32 *) pCur->knn_data->rowids)[pCur->knn_data->current_idx]; i32 rowid = ((i32 *)pCur->knn_data->rowids)[pCur->knn_data->current_idx];
sqlite3_result_blob( sqlite3_result_blob(
context, context, p->blob->p + (rowid * p->blob->dimensions * sizeof(float)),
p->blob->p + (rowid* p->blob->dimensions * sizeof(float)), p->blob->dimensions * sizeof(float), SQLITE_STATIC);
p->blob->dimensions * sizeof(float), sqlite3_result_subtype(context, p->blob->element_type);
SQLITE_STATIC
);
sqlite3_result_subtype(context,p->blob->element_type);
break; break;
} }
} }
@ -5153,7 +5165,6 @@ static int vec_static_blob_entriesColumn(sqlite3_vtab_cursor *cur, sqlite3_conte
} }
} }
static sqlite3_module vec_static_blob_entriesModule = { static sqlite3_module vec_static_blob_entriesModule = {
/* iVersion */ 3, /* iVersion */ 3,
/* xCreate */ vec_static_blob_entriesCreate, /* xCreate */ vec_static_blob_entriesCreate,
@ -5195,7 +5206,9 @@ void dethrone2(int k, f32 *base_distances, i64 *base_rowids, size_t chunk_size,
size_t ptrA = 0; size_t ptrA = 0;
size_t ptrB = 0; size_t ptrB = 0;
for (int i = 0; i < k; i++) { for (int i = 0; i < k; i++) {
if (ptrA < chunk_size && (ptrB >= k || chunk_distances[chunk_top_idx[ptrA]] < base_distances[ptrB])) { if (ptrA < chunk_size &&
(ptrB >= k ||
chunk_distances[chunk_top_idx[ptrA]] < base_distances[ptrB])) {
(*out_rowids)[i] = chunk_rowids[chunk_top_idx[ptrA]]; (*out_rowids)[i] = chunk_rowids[chunk_top_idx[ptrA]];
(*out_distances)[i] = chunk_distances[chunk_top_idx[ptrA]]; (*out_distances)[i] = chunk_distances[chunk_top_idx[ptrA]];
ptrA++; ptrA++;
@ -5207,13 +5220,12 @@ void dethrone2(int k, f32 *base_distances, i64 *base_rowids, size_t chunk_size,
} }
} }
typedef struct vec_expo_vtab vec_expo_vtab; typedef struct vec_expo_vtab vec_expo_vtab;
struct vec_expo_vtab { struct vec_expo_vtab {
sqlite3_vtab base; sqlite3_vtab base;
sqlite3 * db; sqlite3 *db;
char * table; char *table;
char * column; char *column;
}; };
typedef struct vec_expo_cursor vec_expo_cursor; typedef struct vec_expo_cursor vec_expo_cursor;
@ -5221,15 +5233,14 @@ struct vec_expo_cursor {
sqlite3_vtab_cursor base; sqlite3_vtab_cursor base;
sqlite3_int64 iRowid; sqlite3_int64 iRowid;
vec_sbe_query_plan query_plan; vec_sbe_query_plan query_plan;
struct vec0_query_knn_data * knn_data; struct vec0_query_knn_data *knn_data;
}; };
static int vec_expoConnect(sqlite3 *db, void *pAux, int argc, static int vec_expoConnect(sqlite3 *db, void *pAux, int argc,
const char *const *argv, sqlite3_vtab **ppVtab, const char *const *argv, sqlite3_vtab **ppVtab,
char **pzErr) { char **pzErr) {
vec_expo_vtab *pNew; vec_expo_vtab *pNew;
assert(argc==5); assert(argc == 5);
#define VEC_EXPO_VECTOR 0 #define VEC_EXPO_VECTOR 0
#define VEC_EXPO_DISTANCE 1 #define VEC_EXPO_DISTANCE 1
@ -5252,7 +5263,7 @@ static int vec_expoCreate(sqlite3 *db, void *pAux, int argc,
const char *const *argv, sqlite3_vtab **ppVtab, const char *const *argv, sqlite3_vtab **ppVtab,
char **pzErr) { char **pzErr) {
return vec_expoConnect(db, pAux, argc, argv, ppVtab, pzErr); return vec_expoConnect(db, pAux, argc, argv, ppVtab, pzErr);
} }
static int vec_expoDisconnect(sqlite3_vtab *pVtab) { static int vec_expoDisconnect(sqlite3_vtab *pVtab) {
vec_expo_vtab *p = (vec_expo_vtab *)pVtab; vec_expo_vtab *p = (vec_expo_vtab *)pVtab;
@ -5276,7 +5287,8 @@ static int vec_expoClose(sqlite3_vtab_cursor *cur) {
return SQLITE_OK; return SQLITE_OK;
} }
static int vec_expoBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { static int vec_expoBestIndex(sqlite3_vtab *pVTab,
sqlite3_index_info *pIdxInfo) {
vec_expo_vtab *p = (vec_expo_vtab *)pVTab; vec_expo_vtab *p = (vec_expo_vtab *)pVTab;
int iMatchTerm = -1; int iMatchTerm = -1;
int iLimitTerm = -1; int iLimitTerm = -1;
@ -5303,7 +5315,7 @@ static int vec_expoBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo)
iKTerm = i; iKTerm = i;
} }
} }
if(iMatchTerm >= 0) { if (iMatchTerm >= 0) {
if (iLimitTerm < 0 && iKTerm < 0) { if (iLimitTerm < 0 && iKTerm < 0) {
// TODO: error, match on vector1 should require a limit for KNN // TODO: error, match on vector1 should require a limit for KNN
return SQLITE_ERROR; return SQLITE_ERROR;
@ -5332,7 +5344,7 @@ static int vec_expoBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo)
pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN; pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_KNN;
pIdxInfo->estimatedCost = (double)10; // TODO vtab_value(?) as hint? pIdxInfo->estimatedCost = (double)10; // TODO vtab_value(?) as hint?
pIdxInfo->estimatedRows = 10;// TODO vtab_value(?) as hint? pIdxInfo->estimatedRows = 10; // TODO vtab_value(?) as hint?
pIdxInfo->orderByConsumed = 1; pIdxInfo->orderByConsumed = 1;
pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1; pIdxInfo->aConstraintUsage[iMatchTerm].argvIndex = 1;
@ -5345,8 +5357,7 @@ static int vec_expoBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo)
pIdxInfo->aConstraintUsage[iKTerm].omit = 1; pIdxInfo->aConstraintUsage[iKTerm].omit = 1;
} }
} } else {
else {
pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN; pIdxInfo->idxNum = VEC_SBE__QUERYPLAN_FULLSCAN;
pIdxInfo->estimatedCost = 10000.0; pIdxInfo->estimatedCost = 10000.0;
pIdxInfo->estimatedRows = 10000; pIdxInfo->estimatedRows = 10000;
@ -5359,7 +5370,7 @@ static int vec_expoFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
vec_expo_cursor *pCur = (vec_expo_cursor *)pVtabCursor; vec_expo_cursor *pCur = (vec_expo_cursor *)pVtabCursor;
vec_expo_vtab *p = (vec_expo_vtab *)pCur->base.pVtab; vec_expo_vtab *p = (vec_expo_vtab *)pCur->base.pVtab;
if(idxNum == VEC_SBE__QUERYPLAN_KNN) { if (idxNum == VEC_SBE__QUERYPLAN_KNN) {
pCur->query_plan = VEC_SBE__QUERYPLAN_KNN; pCur->query_plan = VEC_SBE__QUERYPLAN_KNN;
struct vec0_query_knn_data *knn_data = struct vec0_query_knn_data *knn_data =
sqlite3_malloc(sizeof(struct vec0_query_knn_data)); sqlite3_malloc(sizeof(struct vec0_query_knn_data));
@ -5391,17 +5402,18 @@ static int vec_expoFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
f32 *topk_distances = sqlite3_malloc(k * sizeof(f32)); f32 *topk_distances = sqlite3_malloc(k * sizeof(f32));
todo_assert(topk_distances); todo_assert(topk_distances);
sqlite3_stmt * stmtRowids; sqlite3_stmt *stmtRowids;
char * zSql = sqlite3_mprintf("select rowid from \"%w\" ", p->table); char *zSql = sqlite3_mprintf("select rowid from \"%w\" ", p->table);
assert(zSql); assert(zSql);
rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmtRowids, NULL); rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmtRowids, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
sqlite3_blob * baseVectorsBlob; sqlite3_blob *baseVectorsBlob;
sqlite3_blob_open(p->db, "main", p->table, p->column, 1, 0, &baseVectorsBlob); sqlite3_blob_open(p->db, "main", p->table, p->column, 1, 0,
&baseVectorsBlob);
int chunk_size = 200; int chunk_size = 200;
float * chunk = sqlite3_malloc(dimensions * chunk_size * sizeof(float)); float *chunk = sqlite3_malloc(dimensions * chunk_size * sizeof(float));
assert(chunk); assert(chunk);
f32 *chunk_distances = sqlite3_malloc(chunk_size * sizeof(f32)); f32 *chunk_distances = sqlite3_malloc(chunk_size * sizeof(f32));
@ -5412,29 +5424,30 @@ static int vec_expoFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
i64 *chunk_rowids = sqlite3_malloc(chunk_size * sizeof(i64)); i64 *chunk_rowids = sqlite3_malloc(chunk_size * sizeof(i64));
todo_assert(chunk_rowids); todo_assert(chunk_rowids);
while (true) {
while(true) {
int nused = 0; int nused = 0;
for(int i = 0; i < chunk_size; i++) { for (int i = 0; i < chunk_size; i++) {
rc = sqlite3_step(stmtRowids); rc = sqlite3_step(stmtRowids);
if(rc == SQLITE_DONE) { if (rc == SQLITE_DONE) {
break; break;
} }
assert(rc == SQLITE_ROW); assert(rc == SQLITE_ROW);
nused = i+1; nused = i + 1;
i64 rowid = sqlite3_column_int64(stmtRowids, 0); i64 rowid = sqlite3_column_int64(stmtRowids, 0);
chunk_rowids[i] = rowid; chunk_rowids[i] = rowid;
rc = sqlite3_blob_reopen(baseVectorsBlob, rowid); rc = sqlite3_blob_reopen(baseVectorsBlob, rowid);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
assert(sqlite3_blob_bytes(baseVectorsBlob) == dimensions * sizeof(float)); assert(sqlite3_blob_bytes(baseVectorsBlob) ==
sqlite3_blob_read(baseVectorsBlob, &chunk[i * dimensions], dimensions * sizeof(float), 0); dimensions * sizeof(float));
sqlite3_blob_read(baseVectorsBlob, &chunk[i * dimensions],
dimensions * sizeof(float), 0);
} }
for(int i = 0; i < nused; i++) { for (int i = 0; i < nused; i++) {
const f32 *base_i = (chunk) + (i * dimensions); const f32 *base_i = (chunk) + (i * dimensions);
chunk_distances[i] = distance_l2_sqr_float(base_i, (f32 *)queryVector, &dimensions); chunk_distances[i] =
distance_l2_sqr_float(base_i, (f32 *)queryVector, &dimensions);
} }
i32 *chunk_top_idxs = sqlite3_malloc(nused * sizeof(i32)); i32 *chunk_top_idxs = sqlite3_malloc(nused * sizeof(i32));
@ -5443,8 +5456,8 @@ static int vec_expoFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
i64 *out_rowids; i64 *out_rowids;
f32 *out_distances; f32 *out_distances;
dethrone2(k, topk_distances, topk_rowids, /*chunk_size*/ nused, chunk_top_idxs, dethrone2(k, topk_distances, topk_rowids, /*chunk_size*/ nused,
chunk_distances, chunk_rowids, chunk_top_idxs, chunk_distances, chunk_rowids,
&out_rowids, &out_distances); &out_rowids, &out_distances);
for (int i = 0; i < k; i++) { for (int i = 0; i < k; i++) {
@ -5455,21 +5468,20 @@ static int vec_expoFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
sqlite3_free(out_distances); sqlite3_free(out_distances);
sqlite3_free(chunk_top_idxs); sqlite3_free(chunk_top_idxs);
if(nused < chunk_size) break; if (nused < chunk_size)
break;
} }
sqlite3_blob_close(baseVectorsBlob); sqlite3_blob_close(baseVectorsBlob);
sqlite3_finalize(stmtRowids); sqlite3_finalize(stmtRowids);
cleanup(queryVector); cleanup(queryVector);
knn_data->current_idx = 0; knn_data->current_idx = 0;
knn_data->k = k; knn_data->k = k;
knn_data->rowids = topk_rowids; knn_data->rowids = topk_rowids;
knn_data->distances = topk_distances; knn_data->distances = topk_distances;
pCur->knn_data = knn_data; pCur->knn_data = knn_data;
} } else {
else {
pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN; pCur->query_plan = VEC_SBE__QUERYPLAN_FULLSCAN;
pCur->iRowid = 0; pCur->iRowid = 0;
} }
@ -5479,7 +5491,7 @@ static int vec_expoFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum,
static int vec_expoRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) { static int vec_expoRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
vec_expo_cursor *pCur = (vec_expo_cursor *)cur; vec_expo_cursor *pCur = (vec_expo_cursor *)cur;
switch(pCur->query_plan) { switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
*pRowid = pCur->iRowid; *pRowid = pCur->iRowid;
break; break;
@ -5495,7 +5507,7 @@ static int vec_expoRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid) {
static int vec_expoNext(sqlite3_vtab_cursor *cur) { static int vec_expoNext(sqlite3_vtab_cursor *cur) {
vec_expo_cursor *pCur = (vec_expo_cursor *)cur; vec_expo_cursor *pCur = (vec_expo_cursor *)cur;
switch(pCur->query_plan) { switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
pCur->iRowid++; pCur->iRowid++;
return SQLITE_OK; return SQLITE_OK;
@ -5505,21 +5517,19 @@ static int vec_expoNext(sqlite3_vtab_cursor *cur) {
return SQLITE_OK; return SQLITE_OK;
} }
} }
} }
static int vec_expoEof(sqlite3_vtab_cursor *cur) { static int vec_expoEof(sqlite3_vtab_cursor *cur) {
vec_expo_cursor *pCur = (vec_expo_cursor *)cur; vec_expo_cursor *pCur = (vec_expo_cursor *)cur;
vec_expo_vtab * p = (vec_expo_vtab *) pCur->base.pVtab; vec_expo_vtab *p = (vec_expo_vtab *)pCur->base.pVtab;
switch(pCur->query_plan) { switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
return 1;//(size_t) pCur->iRowid >= p->blob->nvectors; return 1; //(size_t) pCur->iRowid >= p->blob->nvectors;
} }
case VEC_SBE__QUERYPLAN_KNN: { case VEC_SBE__QUERYPLAN_KNN: {
return pCur->knn_data->current_idx >= pCur->knn_data->k; return pCur->knn_data->current_idx >= pCur->knn_data->k;
} }
} }
} }
static int vec_expoColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context, static int vec_expoColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context,
@ -5527,17 +5537,18 @@ static int vec_expoColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context,
vec_expo_cursor *pCur = (vec_expo_cursor *)cur; vec_expo_cursor *pCur = (vec_expo_cursor *)cur;
vec_expo_vtab *p = (vec_expo_vtab *)cur->pVtab; vec_expo_vtab *p = (vec_expo_vtab *)cur->pVtab;
switch(pCur->query_plan) { switch (pCur->query_plan) {
case VEC_SBE__QUERYPLAN_FULLSCAN: { case VEC_SBE__QUERYPLAN_FULLSCAN: {
return SQLITE_OK; return SQLITE_OK;
} }
case VEC_SBE__QUERYPLAN_KNN: { case VEC_SBE__QUERYPLAN_KNN: {
switch(i) { switch (i) {
case VEC_EXPO_VECTOR: { case VEC_EXPO_VECTOR: {
break; break;
} }
case VEC_EXPO_DISTANCE: { case VEC_EXPO_DISTANCE: {
sqlite3_result_double(context, pCur->knn_data->distances[pCur->knn_data->current_idx]); sqlite3_result_double(
context, pCur->knn_data->distances[pCur->knn_data->current_idx]);
break; break;
} }
} }
@ -5546,7 +5557,6 @@ static int vec_expoColumn(sqlite3_vtab_cursor *cur, sqlite3_context *context,
} }
} }
static sqlite3_module vec_expoModule = { static sqlite3_module vec_expoModule = {
/* iVersion */ 3, /* iVersion */ 3,
/* xCreate */ vec_expoCreate, /* xCreate */ vec_expoCreate,
@ -5689,7 +5699,7 @@ __declspec(dllexport)
#define SQLITE_RESULT_SUBTYPE 0x001000000 #define SQLITE_RESULT_SUBTYPE 0x001000000
#endif #endif
int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg, int sqlite3_vec_init(sqlite3 *db, char **pzErrMsg,
const sqlite3_api_routines *pApi) { const sqlite3_api_routines *pApi) {
SQLITE_EXTENSION_INIT2(pApi); SQLITE_EXTENSION_INIT2(pApi);
int rc = SQLITE_OK; int rc = SQLITE_OK;
@ -5727,18 +5737,18 @@ __declspec(dllexport)
// clang-format on // clang-format on
}; };
#ifdef SQLITE_VEC_ENABLE_EXPERIMENTAL #ifdef SQLITE_VEC_ENABLE_EXPERIMENTAL
vec_static_blob_data * static_blob_data; vec_static_blob_data *static_blob_data;
static_blob_data = sqlite3_malloc(sizeof(*static_blob_data)); static_blob_data = sqlite3_malloc(sizeof(*static_blob_data));
todo_assert(static_blob_data); todo_assert(static_blob_data);
memset(static_blob_data, 0, sizeof(*static_blob_data)); memset(static_blob_data, 0, sizeof(*static_blob_data));
#endif #endif
static const struct { static const struct {
char *name; char *name;
const sqlite3_module *module; const sqlite3_module *module;
void * p; void *p;
void(*xDestroy)(void*); void (*xDestroy)(void *);
} aMod[] = { } aMod[] = {
// clang-format off // clang-format off
{"vec0", &vec0Module, NULL, NULL}, {"vec0", &vec0Module, NULL, NULL},
@ -5767,19 +5777,22 @@ __declspec(dllexport)
return rc; return rc;
} }
} }
#ifdef SQLITE_VEC_ENABLE_EXPERIMENTAL #ifdef SQLITE_VEC_ENABLE_EXPERIMENTAL
rc = sqlite3_create_module_v2(db, "vec_static_blobs", &vec_static_blobsModule, static_blob_data, sqlite3_free); rc = sqlite3_create_module_v2(db, "vec_static_blobs", &vec_static_blobsModule,
static_blob_data, sqlite3_free);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
rc = sqlite3_create_module_v2(db, "vec_static_blob_entries", &vec_static_blob_entriesModule, static_blob_data, NULL); rc = sqlite3_create_module_v2(db, "vec_static_blob_entries",
&vec_static_blob_entriesModule,
static_blob_data, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
rc = sqlite3_create_module_v2(db, "vec_expo", &vec_expoModule, NULL, NULL); rc = sqlite3_create_module_v2(db, "vec_expo", &vec_expoModule, NULL, NULL);
assert(rc == SQLITE_OK); assert(rc == SQLITE_OK);
#endif #endif
return SQLITE_OK; return SQLITE_OK;
} }
int sqlite3_vec_fs_read_init(sqlite3 *db, char **pzErrMsg, int sqlite3_vec_fs_read_init(sqlite3 *db, char **pzErrMsg,
const sqlite3_api_routines *pApi) { const sqlite3_api_routines *pApi) {
UNUSED_PARAMETER(pzErrMsg); UNUSED_PARAMETER(pzErrMsg);
SQLITE_EXTENSION_INIT2(pApi); SQLITE_EXTENSION_INIT2(pApi);