mirror of
https://github.com/asg017/sqlite-vec.git
synced 2026-04-25 00:36:56 +02:00
Merge branch 'main' of github.com:asg017/sqlite-vec
This commit is contained in:
Alex Garcia
commit
fb31b780f5
3 changed files with 492 additions and 2 deletions
139
sqlite-vec.c
139
sqlite-vec.c
|
|
@ -5305,11 +5305,21 @@ static int vec0Close(sqlite3_vtab_cursor *cur) {
|
|||
typedef enum {
|
||||
// If any values are updated, please update the ARCHITECTURE.md docs accordingly!
|
||||
|
||||
// ~~~ KNN QUERIES ~~~ //
|
||||
VEC0_IDXSTR_KIND_KNN_MATCH = '{',
|
||||
VEC0_IDXSTR_KIND_KNN_K = '}',
|
||||
VEC0_IDXSTR_KIND_KNN_ROWID_IN = '[',
|
||||
// argv[i] is a constraint on a PARTITON KEY column in a KNN query
|
||||
//
|
||||
VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT = ']',
|
||||
|
||||
// argv[i] is a constraint on the distance column in a KNN query
|
||||
VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT = '*',
|
||||
|
||||
// ~~~ POINT QUERIES ~~~ //
|
||||
VEC0_IDXSTR_KIND_POINT_ID = '!',
|
||||
|
||||
// ~~~ ??? ~~~ //
|
||||
VEC0_IDXSTR_KIND_METADATA_CONSTRAINT = '&',
|
||||
} vec0_idxstr_kind;
|
||||
|
||||
|
|
@ -5318,11 +5328,22 @@ typedef enum {
|
|||
typedef enum {
|
||||
// If any values are updated, please update the ARCHITECTURE.md docs accordingly!
|
||||
|
||||
// Equality constraint on a PARTITON KEY column, ex `user_id = 123`
|
||||
VEC0_PARTITION_OPERATOR_EQ = 'a',
|
||||
|
||||
// "Greater than" constraint on a PARTITON KEY column, ex `year > 2024`
|
||||
VEC0_PARTITION_OPERATOR_GT = 'b',
|
||||
|
||||
// "Less than or equal to" constraint on a PARTITON KEY column, ex `year <= 2024`
|
||||
VEC0_PARTITION_OPERATOR_LE = 'c',
|
||||
|
||||
// "Less than" constraint on a PARTITON KEY column, ex `year < 2024`
|
||||
VEC0_PARTITION_OPERATOR_LT = 'd',
|
||||
|
||||
// "Greater than or equal to" constraint on a PARTITON KEY column, ex `year >= 2024`
|
||||
VEC0_PARTITION_OPERATOR_GE = 'e',
|
||||
|
||||
// "Not equal to" constraint on a PARTITON KEY column, ex `year != 2024`
|
||||
VEC0_PARTITION_OPERATOR_NE = 'f',
|
||||
} vec0_partition_operator;
|
||||
typedef enum {
|
||||
|
|
@ -5335,6 +5356,15 @@ typedef enum {
|
|||
VEC0_METADATA_OPERATOR_IN = 'g',
|
||||
} vec0_metadata_operator;
|
||||
|
||||
|
||||
typedef enum {
|
||||
|
||||
VEC0_DISTANCE_CONSTRAINT_GT = 'a',
|
||||
VEC0_DISTANCE_CONSTRAINT_GE = 'b',
|
||||
VEC0_DISTANCE_CONSTRAINT_LT = 'c',
|
||||
VEC0_DISTANCE_CONSTRAINT_LE = 'd',
|
||||
} vec0_distance_constraint_operator;
|
||||
|
||||
static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
|
||||
vec0_vtab *p = (vec0_vtab *)pVTab;
|
||||
/**
|
||||
|
|
@ -5494,6 +5524,7 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
|
|||
}
|
||||
#endif
|
||||
|
||||
// find any PARTITION KEY column constraints
|
||||
for (int i = 0; i < pIdxInfo->nConstraint; i++) {
|
||||
if (!pIdxInfo->aConstraint[i].usable)
|
||||
continue;
|
||||
|
|
@ -5548,6 +5579,7 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
|
|||
|
||||
}
|
||||
|
||||
// find any metadata column constraints
|
||||
for (int i = 0; i < pIdxInfo->nConstraint; i++) {
|
||||
if (!pIdxInfo->aConstraint[i].usable)
|
||||
continue;
|
||||
|
|
@ -5644,6 +5676,58 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
|
|||
|
||||
}
|
||||
|
||||
// find any distance column constraints
|
||||
for (int i = 0; i < pIdxInfo->nConstraint; i++) {
|
||||
if (!pIdxInfo->aConstraint[i].usable)
|
||||
continue;
|
||||
|
||||
int iColumn = pIdxInfo->aConstraint[i].iColumn;
|
||||
int op = pIdxInfo->aConstraint[i].op;
|
||||
if(op == SQLITE_INDEX_CONSTRAINT_LIMIT || op == SQLITE_INDEX_CONSTRAINT_OFFSET) {
|
||||
continue;
|
||||
}
|
||||
if(vec0_column_distance_idx(p) != iColumn) {
|
||||
continue;
|
||||
}
|
||||
|
||||
char value = 0;
|
||||
switch(op) {
|
||||
case SQLITE_INDEX_CONSTRAINT_GT: {
|
||||
value = VEC0_DISTANCE_CONSTRAINT_GT;
|
||||
break;
|
||||
}
|
||||
case SQLITE_INDEX_CONSTRAINT_GE: {
|
||||
value = VEC0_DISTANCE_CONSTRAINT_GE;
|
||||
break;
|
||||
}
|
||||
case SQLITE_INDEX_CONSTRAINT_LT: {
|
||||
value = VEC0_DISTANCE_CONSTRAINT_LT;
|
||||
break;
|
||||
}
|
||||
case SQLITE_INDEX_CONSTRAINT_LE: {
|
||||
value = VEC0_DISTANCE_CONSTRAINT_LE;
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
// IMP TODO
|
||||
rc = SQLITE_ERROR;
|
||||
vtab_set_error(
|
||||
pVTab,
|
||||
"Illegal WHERE constraint on distance column in a KNN query. "
|
||||
"Only one of GT, GE, LT, LE constraints are allowed."
|
||||
);
|
||||
goto done;
|
||||
}
|
||||
}
|
||||
|
||||
pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++;
|
||||
pIdxInfo->aConstraintUsage[i].omit = 1;
|
||||
sqlite3_str_appendchar(idxStr, 1, VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT);
|
||||
sqlite3_str_appendchar(idxStr, 1, value);
|
||||
sqlite3_str_appendchar(idxStr, 1, '_');
|
||||
sqlite3_str_appendchar(idxStr, 1, '_');
|
||||
}
|
||||
|
||||
|
||||
|
||||
pIdxInfo->idxNum = iMatchVectorTerm;
|
||||
|
|
@ -5672,7 +5756,6 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
|
|||
}
|
||||
pIdxInfo->needToFreeIdxStr = 1;
|
||||
|
||||
|
||||
rc = SQLITE_OK;
|
||||
|
||||
done:
|
||||
|
|
@ -6560,12 +6643,15 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
|
|||
int numValueEntries = (idxStrLength-1) / 4;
|
||||
assert(numValueEntries == argc);
|
||||
int hasMetadataFilters = 0;
|
||||
int hasDistanceConstraints = 0;
|
||||
for(int i = 0; i < argc; i++) {
|
||||
int idx = 1 + (i * 4);
|
||||
char kind = idxStr[idx + 0];
|
||||
if(kind == VEC0_IDXSTR_KIND_METADATA_CONSTRAINT) {
|
||||
hasMetadataFilters = 1;
|
||||
break;
|
||||
}
|
||||
else if(kind == VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT) {
|
||||
hasDistanceConstraints = 1;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -6752,6 +6838,55 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
|
|||
chunk_distances[i] = result;
|
||||
}
|
||||
|
||||
if(hasDistanceConstraints) {
|
||||
for(int i = 0; i < argc; i++) {
|
||||
int idx = 1 + (i * 4);
|
||||
char kind = idxStr[idx + 0];
|
||||
// TODO casts f64 to f32, is that a problem?
|
||||
f32 target = (f32) sqlite3_value_double(argv[i]);
|
||||
|
||||
if(kind != VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT) {
|
||||
continue;
|
||||
}
|
||||
vec0_distance_constraint_operator op = idxStr[idx + 1];
|
||||
|
||||
switch(op) {
|
||||
case VEC0_DISTANCE_CONSTRAINT_GE: {
|
||||
for(int i = 0; i < p->chunk_size;i++) {
|
||||
if(bitmap_get(b, i) && !(chunk_distances[i] >= target)) {
|
||||
bitmap_set(b, i, 0);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case VEC0_DISTANCE_CONSTRAINT_GT: {
|
||||
for(int i = 0; i < p->chunk_size;i++) {
|
||||
if(bitmap_get(b, i) && !(chunk_distances[i] > target)) {
|
||||
bitmap_set(b, i, 0);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case VEC0_DISTANCE_CONSTRAINT_LE: {
|
||||
for(int i = 0; i < p->chunk_size;i++) {
|
||||
if(bitmap_get(b, i) && !(chunk_distances[i] <= target)) {
|
||||
bitmap_set(b, i, 0);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case VEC0_DISTANCE_CONSTRAINT_LT: {
|
||||
for(int i = 0; i < p->chunk_size;i++) {
|
||||
if(bitmap_get(b, i) && !(chunk_distances[i] < target)) {
|
||||
bitmap_set(b, i, 0);
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int used1;
|
||||
min_idx(chunk_distances, p->chunk_size, b, chunk_topk_idxs,
|
||||
min(k, p->chunk_size), bTaken, &used1);
|
||||
|
|
|
|||
273
tests/__snapshots__/test-knn-distance-constraints.ambr
Normal file
273
tests/__snapshots__/test-knn-distance-constraints.ambr
Normal file
|
|
@ -0,0 +1,273 @@
|
|||
# serializer version: 1
|
||||
# name: test_normal
|
||||
OrderedDict({
|
||||
'sql': 'SELECT * FROM v',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 1,
|
||||
'embedding': b'\x00\x00\x80?',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 2,
|
||||
'embedding': b'\x00\x00\x00@',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 3,
|
||||
'embedding': b'\x00\x00@@',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 4,
|
||||
'embedding': b'\x00\x00\x80@',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 5,
|
||||
'embedding': b'\x00\x00\xa0@',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 6,
|
||||
'embedding': b'\x00\x00\xc0@',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 7,
|
||||
'embedding': b'\x00\x00\xe0@',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 8,
|
||||
'embedding': b'\x00\x00\x00A',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 9,
|
||||
'embedding': b'\x00\x00\x10A',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 10,
|
||||
'embedding': b'\x00\x00 A',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 11,
|
||||
'embedding': b'\x00\x000A',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 12,
|
||||
'embedding': b'\x00\x00@A',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 13,
|
||||
'embedding': b'\x00\x00PA',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 14,
|
||||
'embedding': b'\x00\x00`A',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 15,
|
||||
'embedding': b'\x00\x00pA',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 16,
|
||||
'embedding': b'\x00\x00\x80A',
|
||||
'is_odd': 0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 17,
|
||||
'embedding': b'\x00\x00\x88A',
|
||||
'is_odd': 1,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.1
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? ',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 1,
|
||||
'distance': 0.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 2,
|
||||
'distance': 1.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 3,
|
||||
'distance': 2.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 4,
|
||||
'distance': 3.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 5,
|
||||
'distance': 4.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.2
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND distance > 5',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 7,
|
||||
'distance': 6.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 8,
|
||||
'distance': 7.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 9,
|
||||
'distance': 8.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 10,
|
||||
'distance': 9.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 11,
|
||||
'distance': 10.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.3
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND distance >= 5',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 6,
|
||||
'distance': 5.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 7,
|
||||
'distance': 6.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 8,
|
||||
'distance': 7.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 9,
|
||||
'distance': 8.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 10,
|
||||
'distance': 9.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.4
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND distance < 3',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 1,
|
||||
'distance': 0.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 2,
|
||||
'distance': 1.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 3,
|
||||
'distance': 2.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.5
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND distance <= 3',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 1,
|
||||
'distance': 0.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 2,
|
||||
'distance': 1.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 3,
|
||||
'distance': 2.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 4,
|
||||
'distance': 3.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.6
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND distance > 7 AND distance <= 10',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 9,
|
||||
'distance': 8.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 10,
|
||||
'distance': 9.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 11,
|
||||
'distance': 10.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.7
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND distance BETWEEN 7 AND 10',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 8,
|
||||
'distance': 7.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 9,
|
||||
'distance': 8.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 10,
|
||||
'distance': 9.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 11,
|
||||
'distance': 10.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
# name: test_normal.8
|
||||
OrderedDict({
|
||||
'sql': 'select rowid, distance from v where embedding match ? and k = ? AND is_odd == TRUE AND distance BETWEEN 7 AND 10',
|
||||
'rows': list([
|
||||
OrderedDict({
|
||||
'rowid': 9,
|
||||
'distance': 8.0,
|
||||
}),
|
||||
OrderedDict({
|
||||
'rowid': 11,
|
||||
'distance': 10.0,
|
||||
}),
|
||||
]),
|
||||
})
|
||||
# ---
|
||||
82
tests/test-knn-distance-constraints.py
Normal file
82
tests/test-knn-distance-constraints.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
import sqlite3
|
||||
from collections import OrderedDict
|
||||
|
||||
|
||||
def test_normal(db, snapshot):
|
||||
db.execute("create virtual table v using vec0(embedding float[1], is_odd boolean, chunk_size=8)")
|
||||
db.executemany(
|
||||
"insert into v(rowid, is_odd, embedding) values (?1, ?1 % 2, ?2)",
|
||||
[
|
||||
[1, "[1]"],
|
||||
[2, "[2]"],
|
||||
[3, "[3]"],
|
||||
[4, "[4]"],
|
||||
[5, "[5]"],
|
||||
[6, "[6]"],
|
||||
[7, "[7]"],
|
||||
[8, "[8]"],
|
||||
[9, "[9]"],
|
||||
[10, "[10]"],
|
||||
[11, "[11]"],
|
||||
[12, "[12]"],
|
||||
[13, "[13]"],
|
||||
[14, "[14]"],
|
||||
[15, "[15]"],
|
||||
[16, "[16]"],
|
||||
[17, "[17]"],
|
||||
],
|
||||
)
|
||||
assert exec(db,"SELECT * FROM v") == snapshot()
|
||||
|
||||
BASE_KNN = "select rowid, distance from v where embedding match ? and k = ? "
|
||||
assert exec(db, BASE_KNN, ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND distance > 5", ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND distance >= 5", ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND distance < 3", ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND distance <= 3", ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND distance > 7 AND distance <= 10", ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND distance BETWEEN 7 AND 10", ["[1]", 5]) == snapshot()
|
||||
assert exec(db, BASE_KNN + "AND is_odd == TRUE AND distance BETWEEN 7 AND 10", ["[1]", 5]) == snapshot()
|
||||
|
||||
|
||||
class Row:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return repr()
|
||||
|
||||
|
||||
def exec(db, sql, parameters=[]):
|
||||
try:
|
||||
rows = db.execute(sql, parameters).fetchall()
|
||||
except (sqlite3.OperationalError, sqlite3.DatabaseError) as e:
|
||||
return {
|
||||
"error": e.__class__.__name__,
|
||||
"message": str(e),
|
||||
}
|
||||
a = []
|
||||
for row in rows:
|
||||
o = OrderedDict()
|
||||
for k in row.keys():
|
||||
o[k] = row[k]
|
||||
a.append(o)
|
||||
result = OrderedDict()
|
||||
result["sql"] = sql
|
||||
result["rows"] = a
|
||||
return result
|
||||
|
||||
|
||||
def vec0_shadow_table_contents(db, v):
|
||||
shadow_tables = [
|
||||
row[0]
|
||||
for row in db.execute(
|
||||
"select name from sqlite_master where name like ? order by 1", [f"{v}_%"]
|
||||
).fetchall()
|
||||
]
|
||||
o = {}
|
||||
for shadow_table in shadow_tables:
|
||||
if shadow_table.endswith("_info"):
|
||||
continue
|
||||
o[shadow_table] = exec(db, f"select * from {shadow_table}")
|
||||
return o
|
||||
Loading…
Add table
Add a link
Reference in a new issue