limit checks

This commit is contained in:
Alex Garcia 2024-08-01 02:45:51 -07:00
parent 09c3f05759
commit e379c205c8
2 changed files with 32 additions and 1 deletions

View file

@ -1391,7 +1391,7 @@ static void vec_quantize_int8(sqlite3_context *context, int argc,
(sqlite3_stricmp((const char *)sqlite3_value_text(argv[1]), "unit") != (sqlite3_stricmp((const char *)sqlite3_value_text(argv[1]), "unit") !=
0)) { 0)) {
sqlite3_result_error( sqlite3_result_error(
context, "2nd argument to vec_quantize_i8() must be 'unit'.", -1); context, "2nd argument to vec_quantize_int8() must be 'unit'.", -1);
sqlite3_free(out); sqlite3_free(out);
goto cleanup; goto cleanup;
} }
@ -3939,6 +3939,15 @@ static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv,
VEC0_MAX_VECTOR_COLUMNS); VEC0_MAX_VECTOR_COLUMNS);
goto error; goto error;
} }
#define SQLITE_VEC_VEC0_MAX_DIMENSIONS 8192
if(c.dimensions > SQLITE_VEC_VEC0_MAX_DIMENSIONS) {
sqlite3_free(c.name);
*pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
"Dimension on vector column too large, provided %lld, maximum %lld",
(i64) c.dimensions,
SQLITE_VEC_VEC0_MAX_DIMENSIONS);
goto error;
}
memcpy(&pNew->vector_columns[numVectorColumns], &c, sizeof(c)); memcpy(&pNew->vector_columns[numVectorColumns], &c, sizeof(c));
numVectorColumns++; numVectorColumns++;
continue; continue;
@ -3990,6 +3999,12 @@ static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv,
"chunk_size must be divisible by 8"); "chunk_size must be divisible by 8");
goto error; goto error;
} }
#define SQLITE_VEC_CHUNK_SIZE_MAX 4096
if (chunk_size > SQLITE_VEC_CHUNK_SIZE_MAX) {
*pzErr = sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR
"chunk_size too large");
goto error;
}
} else { } else {
// IMP: V27642_11712 // IMP: V27642_11712
*pzErr = sqlite3_mprintf( *pzErr = sqlite3_mprintf(
@ -4907,6 +4922,13 @@ int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum,
rc = SQLITE_ERROR; rc = SQLITE_ERROR;
goto cleanup; goto cleanup;
} }
#define SQLITE_VEC_VEC0_K_MAX 4096
if(k > SQLITE_VEC_VEC0_K_MAX) {
vtab_set_error(
&p->base, "k value in knn query too large, provided %lld and the limit is %lld", k, SQLITE_VEC_VEC0_K_MAX);
rc = SQLITE_ERROR;
goto cleanup;
}
if (k == 0) { if (k == 0) {
knn_data->k = 0; knn_data->k = 0;

View file

@ -264,7 +264,16 @@ def test_vec_static_blob_entries():
"v": "[0.300000,0.300000,0.300000,0.300000]", "v": "[0.300000,0.300000,0.300000,0.300000]",
}, },
] ]
def test_limits():
db = connect(EXT_PATH)
with _raises("vec0 constructor error: Dimension on vector column too large, provided 8193, maximum 8192"):
db.execute("create virtual table v using vec0(a float[8193])")
with _raises("vec0 constructor error: chunk_size too large"):
db.execute("create virtual table v using vec0(a float[4], chunk_size=8200)")
db.execute('create virtual table v using vec0(a float[1])')
with _raises("k value in knn query too large, provided 8193 and the limit is 4096"):
db.execute("select * from v where a match '[0.1]' and k = 8193")
def test_funcs(): def test_funcs():
funcs = list( funcs = list(