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

more snpashot tests

Browse source
This commit is contained in:
Alex Garcia 2024-11-13 21:52:35 -08:00
parent b04b64759c
commit 06e28da3a5
3 changed files with 570 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

26
sqlite-vec.c
View file

@ -4061,7 +4061,13 @@ int vec0_result_metadata_value_for_rowid(vec0_vtab *p, i64 rowid, int metadata_i
} }
switch(p->metadata_columns[metadata_idx].kind) { switch(p->metadata_columns[metadata_idx].kind) {
case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
// TODO handle boolean values u8 block;
rc = sqlite3_blob_read(blobValue, &block, sizeof(block), chunk_offset / CHAR_BIT);
if(rc != SQLITE_OK) {
goto done;
}
int value = block >> ((chunk_offset % CHAR_BIT)) & 1;
sqlite3_result_int(context, value);
break; break;
} }
case VEC0_METADATA_COLUMN_KIND_INT32: { case VEC0_METADATA_COLUMN_KIND_INT32: {
@ -7410,7 +7416,7 @@ int vec0_insert_metadata_values(vec0_vtab *p, int argc, sqlite3_value ** argv, i
// verify input value matches column type // verify input value matches column type
switch(kind) { switch(kind) {
case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
if(sqlite3_value_type(v) != SQLITE_INTEGER || (sqlite3_value_int(v) != 0) || (sqlite3_value_int(v) != 1)) { if(sqlite3_value_type(v) != SQLITE_INTEGER || ((sqlite3_value_int(v) != 0) && (sqlite3_value_int(v) != 1))) {
rc = SQLITE_ERROR; rc = SQLITE_ERROR;
vtab_set_error(&p->base, "Expected 0 or 1 for BOOLEAN metadata column %s.%s.%s", p->schemaName, p->tableName, p->shadowMetadataChunksNames[metadata_idx]); vtab_set_error(&p->base, "Expected 0 or 1 for BOOLEAN metadata column %s.%s.%s", p->schemaName, p->tableName, p->shadowMetadataChunksNames[metadata_idx]);
goto done; goto done;
@ -7447,7 +7453,21 @@ int vec0_insert_metadata_values(vec0_vtab *p, int argc, sqlite3_value ** argv, i
} }
switch(kind) { switch(kind) {
case VEC0_METADATA_COLUMN_KIND_BOOLEAN: { case VEC0_METADATA_COLUMN_KIND_BOOLEAN: {
//sqlite3_blob_write(blobValue, ) u8 block;
int value = sqlite3_value_int(v);
rc = sqlite3_blob_read(blobValue, &block, sizeof(u8), (int) (chunk_offset / CHAR_BIT));
if(rc != SQLITE_OK) {
// TODO
goto done;
}
if (value) {
block |= 1 << (chunk_offset % CHAR_BIT);
} else {
block &= ~(1 << (chunk_offset % CHAR_BIT));
}
rc = sqlite3_blob_write(blobValue, &block, sizeof(u8), chunk_offset / CHAR_BIT);
break; break;
} }
case VEC0_METADATA_COLUMN_KIND_INT32: { case VEC0_METADATA_COLUMN_KIND_INT32: {

412
tests/__snapshots__/test-metadata.ambr
View file

@ -5,6 +5,214 @@
'message': 'vec0 constructor error: More than 16 metadata columns were provided', 'message': 'vec0 constructor error: More than 16 metadata columns were provided',
}) })
# --- # ---
# name: test_idxstr
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and is_favorited = true",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Aa_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-float !=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and mean_rating != NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Df_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-float <=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and mean_rating <= NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Dc_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-float <]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and mean_rating < NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Dd_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-float >=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and mean_rating >= NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&De_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-float >]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and mean_rating > NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Db_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-int !=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and num_reviews != NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Cf_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-int <=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and num_reviews <= NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Cc_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-int <]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and num_reviews < NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Cd_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-int >=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and num_reviews >= NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Ce_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-int >]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and num_reviews > NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Cb_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-text !=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and genre != NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Bf_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-text <=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and genre <= NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Bc_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-text <]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and genre < NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Bd_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-text >=]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and genre >= NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Be_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_idxstr[knn-constraint-text >]
OrderedDict({
'sql': "select * from vec_movies where synopsis_embedding match '' and k = 0 and genre > NULL",
'plan': list([
dict({
'detail': 'SCAN vec_movies VIRTUAL TABLE INDEX 0:3{___}___&Bb_',
'id': 2,
'notused': 0,
'parent': 0,
}),
]),
})
# ---
# name: test_knn.1 # name: test_knn.1
dict({ dict({
'error': 'OperationalError', 'error': 'OperationalError',
@ -466,6 +674,27 @@
}), }),
'vec_movies_metadata_chunks00': OrderedDict({ 'vec_movies_metadata_chunks00': OrderedDict({
'sql': 'select * from vec_movies_metadata_chunks00', 'sql': 'select * from vec_movies_metadata_chunks00',
'rows': list([
OrderedDict({
'rowid': 1,
'data': b'p',
}),
OrderedDict({
'rowid': 2,
'data': b'U',
}),
OrderedDict({
'rowid': 3,
'data': b'\xff',
}),
OrderedDict({
'rowid': 4,
'data': b'\x01',
}),
]),
}),
'vec_movies_metadata_chunks01': OrderedDict({
'sql': 'select * from vec_movies_metadata_chunks01',
'rows': list([ 'rows': list([
OrderedDict({ OrderedDict({
'rowid': 1, 'rowid': 1,
@ -485,8 +714,8 @@
}), }),
]), ]),
}), }),
'vec_movies_metadata_chunks01': OrderedDict({ 'vec_movies_metadata_chunks02': OrderedDict({
'sql': 'select * from vec_movies_metadata_chunks01', 'sql': 'select * from vec_movies_metadata_chunks02',
'rows': list([ 'rows': list([
OrderedDict({ OrderedDict({
'rowid': 1, 'rowid': 1,
@ -506,8 +735,8 @@
}), }),
]), ]),
}), }),
'vec_movies_metadata_chunks02': OrderedDict({ 'vec_movies_metadata_chunks03': OrderedDict({
'sql': 'select * from vec_movies_metadata_chunks02', 'sql': 'select * from vec_movies_metadata_chunks03',
'rows': list([ 'rows': list([
OrderedDict({ OrderedDict({
'rowid': 1, 'rowid': 1,
@ -715,6 +944,7 @@
genre, genre,
num_reviews, num_reviews,
mean_rating, mean_rating,
is_favorited,
distance distance
from vec_movies from vec_movies
where synopsis_embedding match '[15.5]' where synopsis_embedding match '[15.5]'
@ -731,6 +961,7 @@
'genre': 'scifi', 'genre': 'scifi',
'num_reviews': 423, 'num_reviews': 423,
'mean_rating': 4.5, 'mean_rating': 4.5,
'is_favorited': 1,
'distance': 2.5, 'distance': 2.5,
}), }),
OrderedDict({ OrderedDict({
@ -739,6 +970,7 @@
'genre': 'scifi', 'genre': 'scifi',
'num_reviews': 201, 'num_reviews': 201,
'mean_rating': 5.0, 'mean_rating': 5.0,
'is_favorited': 1,
'distance': 2.5, 'distance': 2.5,
}), }),
OrderedDict({ OrderedDict({
@ -747,6 +979,7 @@
'genre': 'scifi', 'genre': 'scifi',
'num_reviews': 342, 'num_reviews': 342,
'mean_rating': 4.0, 'mean_rating': 4.0,
'is_favorited': 1,
'distance': 5.5, 'distance': 5.5,
}), }),
OrderedDict({ OrderedDict({
@ -755,6 +988,7 @@
'genre': 'scifi', 'genre': 'scifi',
'num_reviews': 451, 'num_reviews': 451,
'mean_rating': 4.400000095367432, 'mean_rating': 4.400000095367432,
'is_favorited': 1,
'distance': 6.5, 'distance': 6.5,
}), }),
OrderedDict({ OrderedDict({
@ -763,8 +997,178 @@
'genre': 'scifi', 'genre': 'scifi',
'num_reviews': 301, 'num_reviews': 301,
'mean_rating': 5.0, 'mean_rating': 5.0,
'is_favorited': 0,
'distance': 7.5, 'distance': 7.5,
}), }),
]), ]),
}) })
# --- # ---
# name: test_stress.2
OrderedDict({
'sql': "select movie_id, genre, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and genre = 'horror'",
'rows': list([
OrderedDict({
'movie_id': 24,
'genre': 'horror',
'distance': 76.0,
}),
OrderedDict({
'movie_id': 16,
'genre': 'horror',
'distance': 84.0,
}),
OrderedDict({
'movie_id': 11,
'genre': 'horror',
'distance': 89.0,
}),
OrderedDict({
'movie_id': 6,
'genre': 'horror',
'distance': 94.0,
}),
OrderedDict({
'movie_id': 1,
'genre': 'horror',
'distance': 99.0,
}),
]),
})
# ---
# name: test_stress.3
OrderedDict({
'sql': "select movie_id, genre, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and genre = 'comedy'",
'rows': list([
OrderedDict({
'movie_id': 17,
'genre': 'comedy',
'distance': 83.0,
}),
OrderedDict({
'movie_id': 12,
'genre': 'comedy',
'distance': 88.0,
}),
OrderedDict({
'movie_id': 7,
'genre': 'comedy',
'distance': 93.0,
}),
OrderedDict({
'movie_id': 2,
'genre': 'comedy',
'distance': 98.0,
}),
]),
})
# ---
# name: test_stress.4
OrderedDict({
'sql': "select movie_id, num_reviews, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and num_reviews between 100 and 500",
'rows': list([
OrderedDict({
'movie_id': 25,
'num_reviews': 310,
'distance': 75.0,
}),
OrderedDict({
'movie_id': 24,
'num_reviews': 271,
'distance': 76.0,
}),
OrderedDict({
'movie_id': 22,
'num_reviews': 451,
'distance': 78.0,
}),
OrderedDict({
'movie_id': 21,
'num_reviews': 342,
'distance': 79.0,
}),
OrderedDict({
'movie_id': 20,
'num_reviews': 186,
'distance': 80.0,
}),
]),
})
# ---
# name: test_stress.5
OrderedDict({
'sql': "select movie_id, num_reviews, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and num_reviews >= 500",
'rows': list([
OrderedDict({
'movie_id': 23,
'num_reviews': 522,
'distance': 77.0,
}),
]),
})
# ---
# name: test_stress.6
OrderedDict({
'sql': "select movie_id, mean_rating, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and mean_rating < 3.0",
'rows': list([
OrderedDict({
'movie_id': 19,
'mean_rating': 2.700000047683716,
'distance': 81.0,
}),
OrderedDict({
'movie_id': 12,
'mean_rating': 2.799999952316284,
'distance': 88.0,
}),
OrderedDict({
'movie_id': 7,
'mean_rating': 2.9000000953674316,
'distance': 93.0,
}),
OrderedDict({
'movie_id': 2,
'mean_rating': 2.5999999046325684,
'distance': 98.0,
}),
]),
})
# ---
# name: test_stress.7
OrderedDict({
'sql': "select movie_id, mean_rating, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and mean_rating between 4.0 and 5.0",
'rows': list([
OrderedDict({
'movie_id': 24,
'mean_rating': 4.300000190734863,
'distance': 76.0,
}),
OrderedDict({
'movie_id': 23,
'mean_rating': 4.599999904632568,
'distance': 77.0,
}),
OrderedDict({
'movie_id': 22,
'mean_rating': 4.400000095367432,
'distance': 78.0,
}),
OrderedDict({
'movie_id': 21,
'mean_rating': 4.0,
'distance': 79.0,
}),
OrderedDict({
'movie_id': 20,
'mean_rating': 4.800000190734863,
'distance': 80.0,
}),
]),
})
# ---
# name: test_stress.8
OrderedDict({
'sql': "select movie_id, mean_rating, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and is_favorited = TRUE",
'rows': list([
]),
})
# ---

165
tests/test-metadata.py
View file

@ -80,6 +80,66 @@ def test_knn(db, snapshot):
) )
def test_idxstr(db, snapshot):
db.execute(
"""
create virtual table vec_movies using vec0(
movie_id integer primary key,
synopsis_embedding float[1],
+title text,
is_favorited boolean,
genre text,
num_reviews int,
mean_rating float,
chunk_size=8
);
"""
)
assert (
eqp(
db,
"select * from vec_movies where synopsis_embedding match '' and k = 0 and is_favorited = true",
)
== snapshot()
)
ops = ["<", ">", "<=", ">=", "!="]
for op in ops:
assert eqp(
db,
f"select * from vec_movies where synopsis_embedding match '' and k = 0 and genre {op} NULL",
) == snapshot(name=f"knn-constraint-text {op}")
for op in ops:
assert eqp(
db,
f"select * from vec_movies where synopsis_embedding match '' and k = 0 and num_reviews {op} NULL",
) == snapshot(name=f"knn-constraint-int {op}")
for op in ops:
assert eqp(
db,
f"select * from vec_movies where synopsis_embedding match '' and k = 0 and mean_rating {op} NULL",
) == snapshot(name=f"knn-constraint-float {op}")
# for op in ops:
# assert eqp(
# db,
# f"select * from vec_movies where synopsis_embedding match '' and k = 0 and is_favorited {op} NULL",
# ) == snapshot(name=f"knn-constraint-boolean {op}")
def eqp(db, sql):
o = OrderedDict()
o["sql"] = sql
o["plan"] = [
dict(row) for row in db.execute(f"explain query plan {sql}").fetchall()
]
return o
def test_stress(db, snapshot): def test_stress(db, snapshot):
db.execute( db.execute(
""" """
@ -87,6 +147,7 @@ def test_stress(db, snapshot):
movie_id integer primary key, movie_id integer primary key,
synopsis_embedding float[1], synopsis_embedding float[1],
+title text, +title text,
is_favorited boolean,
genre text, genre text,
num_reviews int, num_reviews int,
mean_rating float, mean_rating float,
@ -97,33 +158,33 @@ def test_stress(db, snapshot):
db.execute( db.execute(
""" """
INSERT INTO vec_movies(movie_id, synopsis_embedding, genre, title, num_reviews, mean_rating) INSERT INTO vec_movies(movie_id, synopsis_embedding, is_favorited, genre, title, num_reviews, mean_rating)
VALUES VALUES
(1, '[1]', 'horror', 'The Conjuring', 153, 4.6), (1, '[1]', 0, 'horror', 'The Conjuring', 153, 4.6),
(2, '[2]', 'comedy', 'Dumb and Dumber', 382, 2.6), (2, '[2]', 0, 'comedy', 'Dumb and Dumber', 382, 2.6),
(3, '[3]', 'scifi', 'Interstellar', 53, 5.0), (3, '[3]', 0, 'scifi', 'Interstellar', 53, 5.0),
(4, '[4]', 'fantasy', 'The Lord of the Rings: The Fellowship of the Ring', 210, 4.2), (4, '[4]', 0, 'fantasy', 'The Lord of the Rings: The Fellowship of the Ring', 210, 4.2),
(5, '[5]', 'documentary', 'An Inconvenient Truth', 93, 3.4), (5, '[5]', 1, 'documentary', 'An Inconvenient Truth', 93, 3.4),
(6, '[6]', 'horror', 'Hereditary', 167, 4.7), (6, '[6]', 1, 'horror', 'Hereditary', 167, 4.7),
(7, '[7]', 'comedy', 'Anchorman: The Legend of Ron Burgundy', 482, 2.9), (7, '[7]', 1, 'comedy', 'Anchorman: The Legend of Ron Burgundy', 482, 2.9),
(8, '[8]', 'scifi', 'Blade Runner 2049', 301, 5.0), (8, '[8]', 0, 'scifi', 'Blade Runner 2049', 301, 5.0),
(9, '[9]', 'fantasy', 'Harry Potter and the Sorcerer''s Stone', 134, 4.1), (9, '[9]', 1, 'fantasy', 'Harry Potter and the Sorcerer''s Stone', 134, 4.1),
(10, '[10]', 'documentary', 'Free Solo', 66, 3.2), (10, '[10]', 0, 'documentary', 'Free Solo', 66, 3.2),
(11, '[11]', 'horror', 'Get Out', 88, 4.9), (11, '[11]', 1, 'horror', 'Get Out', 88, 4.9),
(12, '[12]', 'comedy', 'The Hangover', 59, 2.8), (12, '[12]', 0, 'comedy', 'The Hangover', 59, 2.8),
(13, '[13]', 'scifi', 'The Matrix', 423, 4.5), (13, '[13]', 1, 'scifi', 'The Matrix', 423, 4.5),
(14, '[14]', 'fantasy', 'Pan''s Labyrinth', 275, 3.6), (14, '[14]', 0, 'fantasy', 'Pan''s Labyrinth', 275, 3.6),
(15, '[15]', 'documentary', '13th', 191, 4.4), (15, '[15]', 1, 'documentary', '13th', 191, 4.4),
(16, '[16]', 'horror', 'It Follows', 314, 4.3), (16, '[16]', 0, 'horror', 'It Follows', 314, 4.3),
(17, '[17]', 'comedy', 'Step Brothers', 74, 3.0), (17, '[17]', 1, 'comedy', 'Step Brothers', 74, 3.0),
(18, '[18]', 'scifi', 'Inception', 201, 5.0), (18, '[18]', 1, 'scifi', 'Inception', 201, 5.0),
(19, '[19]', 'fantasy', 'The Shape of Water', 399, 2.7), (19, '[19]', 1, 'fantasy', 'The Shape of Water', 399, 2.7),
(20, '[20]', 'documentary', 'Won''t You Be My Neighbor?', 186, 4.8), (20, '[20]', 1, 'documentary', 'Won''t You Be My Neighbor?', 186, 4.8),
(21, '[21]', 'scifi', 'Gravity', 342, 4.0), (21, '[21]', 1, 'scifi', 'Gravity', 342, 4.0),
(22, '[22]', 'scifi', 'Dune', 451, 4.4), (22, '[22]', 1, 'scifi', 'Dune', 451, 4.4),
(23, '[23]', 'scifi', 'The Martian', 522, 4.6), (23, '[23]', 1, 'scifi', 'The Martian', 522, 4.6),
(24, '[24]', 'horror', 'A Quiet Place', 271, 4.3), (24, '[24]', 1, 'horror', 'A Quiet Place', 271, 4.3),
(25, '[25]', 'fantasy', 'The Chronicles of Narnia: The Lion, the Witch and the Wardrobe', 310, 3.9); (25, '[25]', 1, 'fantasy', 'The Chronicles of Narnia: The Lion, the Witch and the Wardrobe', 310, 3.9);
""" """
) )
@ -139,6 +200,7 @@ def test_stress(db, snapshot):
genre, genre,
num_reviews, num_reviews,
mean_rating, mean_rating,
is_favorited,
distance distance
from vec_movies from vec_movies
where synopsis_embedding match '[15.5]' where synopsis_embedding match '[15.5]'
@ -151,6 +213,57 @@ def test_stress(db, snapshot):
== snapshot() == snapshot()
) )
assert (
exec(
db,
"select movie_id, genre, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and genre = 'horror'",
)
== snapshot()
)
assert (
exec(
db,
"select movie_id, genre, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and genre = 'comedy'",
)
== snapshot()
)
assert (
exec(
db,
"select movie_id, num_reviews, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and num_reviews between 100 and 500",
)
== snapshot()
)
assert (
exec(
db,
"select movie_id, num_reviews, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and num_reviews >= 500",
)
== snapshot()
)
assert (
exec(
db,
"select movie_id, mean_rating, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and mean_rating < 3.0",
)
== snapshot()
)
assert (
exec(
db,
"select movie_id, mean_rating, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and mean_rating between 4.0 and 5.0",
)
== snapshot()
)
assert (
exec(
db,
"select movie_id, mean_rating, distance from vec_movies where synopsis_embedding match '[100]' and k = 5 and is_favorited = TRUE",
)
== snapshot()
)
def exec(db, sql, parameters=[]): def exec(db, sql, parameters=[]):
try: try: