diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index dc9897b6..94ab4575 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -240,20 +240,24 @@ of the `_chunks` query, so the query fails to prepare when the distinct values of its `IN` blocks and its other partition key values together number more than `SQLITE_LIMIT_VARIABLE_NUMBER`. `ISNULL` and `ISNOTNULL` blocks bind nothing. -These conditions compare text with `BINARY` (`LIKE` and `GLOB` keep their own -case rules, and `LIKE` follows `PRAGMA case_sensitive_like`), and they see only -`argv[i]`'s value, not the affinity of the expression that produced it, so -`p IS CAST(5 AS INTEGER)` does not match a text `'5'`. `EQ`, `NE`, and the -range operators are encoded whatever collation the query names. An `IN` or -`IS` constraint is encoded as such only when `sqlite3_vtab_collation()` -reports `BINARY`: an `IN` with another collation is encoded as `EQ`, which -SQLite runs once per value that is distinct under that collation, as before -`vec0` encoded `IN`, and an `IS` with -another collation is left to SQLite. SQLite reports `BINARY` for `IS NOT` -whatever collation the query names, so `xBestIndex` leaves `omit` unset for it -and SQLite re-checks each row. - -The fourth character of the block is a `_` filler. +These conditions see only `argv[i]`'s value, not the affinity of the +expression that produced it, so `p IS CAST(5 AS INTEGER)` does not match a +text `'5'`. On a text key, `EQ` and the range operators compare under the +collation `sqlite3_vtab_collation()` reports, which the fourth character of the +block carries as for a metadata constraint (`'_'`, `'n'` or `'r'`, below) and +`vec0_chunks_iter()` appends to the condition as `COLLATE NOCASE` or +`COLLATE RTRIM`; `xBestIndex` rejects any other collation on them. `LIKE` and +`GLOB` keep their own case rules, and `LIKE` follows +`PRAGMA case_sensitive_like`. An `IN` or `IS` constraint is encoded as such only +when `sqlite3_vtab_collation()` reports `BINARY`: an `IN` with another +collation is encoded as `EQ`, which SQLite runs once per value that is distinct +under that collation, as before `vec0` encoded `IN`, and an `IS` with another +collation is left to SQLite. SQLite reports `BINARY` for `NE` and `IS NOT` +whatever collation the query names, so on a text key `xBestIndex` leaves +`omit` unset for them and SQLite re-checks each row. + +The fourth character of the block is the collation of a text key's `EQ` or +range constraint, and a `_` filler otherwise. #### `VEC0_IDXSTR_KIND_POINT_ID` (`'!'`) @@ -301,6 +305,14 @@ a collation, `LIKE` and `GLOB` have their own case rules, and SQLite reports compares those two bytewise and leaves `omit` unset, so SQLite re-checks each row with the query's collation. +`vec0` keeps text metadata as UTF-8, while SQLite orders text under `BINARY` by +the bytes of the database's encoding, and UTF-16 orders some characters +differently (SQLite registers `BINARY` in each encoding and `NOCASE` and +`RTRIM` for UTF-8 only). So `vec0_init()` reads the encoding with +`PRAGMA encoding` into `textEncoding`, and in a UTF-16 database +`vec0_text_cmp()` compares under `BINARY` by the bytes SQLite stores for each +value, converting its UTF-8 as SQLite's `utf.c` does. + #### `VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT` (`'*'`) `argv[i]` is a constraint on the `distance` column in a KNN query. diff --git a/CHANGELOG.md b/CHANGELOG.md index 326649bf..4aeea076 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### Changed +- KNN `=`, `<`, `<=`, `>`, `>=`, and `BETWEEN` filters on a `text` partition key now compare under the `NOCASE` or `RTRIM` collation the query names, and so does an `IN` that names one, which `vec0` runs as one `=` per value: `name = 'alice' collate nocase` searched no partition keyed `'Alice'`, because `vec0` compared bytewise while telling SQLite the filter needed no second check, and `name collate nocase in ('alice', 'bob')` now also searches the partition keyed `'Alice'`. The range conditions SQLite passes for a `LIKE ... ESCAPE` pattern that begins with literal text compare under `NOCASE` too, so such a query no longer drops keys that differ from the pattern only in case. A filter that names any other collation now fails with `Collation "x" is not supported on the "name" partition key in a vec0 KNN query. Only BINARY, NOCASE, and RTRIM are supported.`, as on a `text` metadata column, and SQLite checks a text key's `!=` filter again, so with a collation it can return fewer than `k` rows. Ported from the unmerged [vlasky#12](https://github.com/vlasky/sqlite-vec/pull/12) ([asg017#329](https://github.com/asg017/sqlite-vec/pull/329)). +- In a database whose encoding is UTF-16, KNN `<`, `<=`, `>`, and `>=` filters on a `text` metadata column under `BINARY` now compare text in the order SQLite gives it there, by the bytes of its UTF-16 form, and so do the range conditions SQLite passes for a `GLOB` or `LIKE` pattern that begins with literal text. They compared the UTF-8 that `vec0` keeps, which orders some characters differently (`U+FFFD` before `U+1F600` in UTF-8 and after it in UTF-16, and `'Ā'` before `'a'` in UTF-16LE alone), so they kept rows a plain scan drops and dropped rows it keeps. `vec0` reads the encoding with `PRAGMA encoding` when it opens a table, and opening now fails with that statement's error when it cannot run, as under an authorizer that denies `SQLITE_PRAGMA`; a build with `SQLITE_OMIT_UTF16`, which has no such pragma, treats every database as UTF-8. `NOCASE` and `RTRIM` filters, which SQLite defines for UTF-8 in every encoding, are unchanged. Ported from the unmerged [vlasky#12](https://github.com/vlasky/sqlite-vec/pull/12) ([asg017#329](https://github.com/asg017/sqlite-vec/pull/329)). - Reading a vector column without a KNN query, as in `select embedding from v`, reads each chunk's vectors through one blob handle instead of opening the chunk's blob again for every row: selecting every vector of a 10,000-row `float[768]` table took 25 ms instead of 2,080 ms, and `select vec_distance_cosine(embedding, ?) from v` 41 ms instead of 2,093 ms. - The linux-arm64 and linux-arm64-musl binaries use NEON, as the darwin-arm64 binary does: `make` enables it on Linux arm64, and `OMIT_SIMD=1` builds without it. An L2 KNN query on a 100,000-row `float[768]` table took 19 ms instead of 57 ms there. Hamming and `int8` L1 distances are unchanged; `float32` L1 distances can differ in the last digit when their double-precision sum rounds, and `float32` L2 distances can differ from 2.1.1's in the last digits, closer to the exact value (at 768 dimensions, a mean relative error of 3.0e-8 instead of 1.7e-7), so a KNN query can order rows with nearly equal distances differently; `int8` L2 distances whose squared differences sum past 2^24 are computed from the exact sum. - Cosine distance between `float32` vectors of more than 16 dimensions, or `int8` vectors of more than 15, uses NEON in NEON builds, which include the darwin-arm64 binary: a KNN query on a 100,000-row `float[768]` table took 26 ms instead of 94 ms on Linux arm64. `float32` distances can differ from the scalar code's in the last digits, as NEON adds the same double-precision products in a different order: differences are rare between random vectors and common between near-duplicate ones, and a vector's distance to itself is between -2.2e-16 and 2.2e-16 in either. `int8` distances match up to 1,024 dimensions; past that, vectors whose squared elements sum past 2^24 can get different distances, usually closer to the exact value, as NEON sums them exactly. Ported from [asg017@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f), except that `float32` sums stay in `double` ([asg017#324](https://github.com/asg017/sqlite-vec/issues/324)), where asg017's kernel sums them in `float32`, and `int8` sums are kept in 64 bits, where asg017's 32-bit sums overflow from 131,072 elements. diff --git a/site/features/vec0.md b/site/features/vec0.md index 8821763c..acae1388 100644 --- a/site/features/vec0.md +++ b/site/features/vec0.md @@ -126,7 +126,9 @@ vec0 KNN query. Only BINARY, NOCASE, and RTRIM are supported.` A `!=` or `IS NOT` condition with `COLLATE` on the value, as in `genre != 'scifi' collate nocase`, can return fewer than `k` rows, because `vec0` compares it bytewise and SQLite drops each returned row that equals the -value under that collation. +value under that collation. `BINARY` orders text as SQLite does in the +database's encoding: in a UTF-16 database, by the bytes of its UTF-16 form, +which puts some characters in another order than UTF-8 does. SQLite does not pass some conditions on metadata columns to `vec0`, and instead applies them to the `k` rows that `vec0` returns. A KNN query with such @@ -159,16 +161,24 @@ condition on a partition key column restricts the search to the partitions whose key satisfies it, except in the forms listed below. `LIKE` and `GLOB` match as in SQLite, and `LIKE` follows `PRAGMA case_sensitive_like`. -When `vec0` uses an `=`, `!=`, `<`, `<=`, `>`, `>=`, or `BETWEEN` condition on -a `text` key, it compares bytewise and ignores any `COLLATE` clause, so -`name = 'alice' collate nocase` does not search a partition keyed `'Alice'`. +An `=`, `<`, `<=`, `>`, `>=`, or `BETWEEN` condition on a `text` key compares +under the `BINARY`, `NOCASE`, or `RTRIM` collation that the query names, so +`name = 'alice' collate nocase` searches a partition keyed `'Alice'`; any other +collation fails the query with `Collation "x" is not supported on the "name" +partition key in a vec0 KNN query. Only BINARY, NOCASE, and RTRIM are +supported.` A `!=` condition with `COLLATE` on the value, as in +`name != 'alice' collate nocase`, can return fewer than `k` rows, because +`vec0` compares it bytewise and SQLite drops each returned row that equals the +value under that collation. + `vec0` uses an `IN` condition that names a collation other than `BINARY`, such -as `name collate nocase in ('alice', 'bob')`, only as one bytewise `=` per -value: it searches the `'alice'` and `'bob'` partitions one at a time and -returns up to `k` rows from each, and the query fails if it uses `LIMIT` -instead of `k = ?`. SQLite first keeps only one of the values that the -collation treats as equal, so `name collate nocase in ('alice', 'Alice')` -searches only one of those two partitions. An `IS NOT` condition with `COLLATE` on the value, as in +as `name collate nocase in ('alice', 'bob')`, only as one `=` per value under +that collation: it searches the partitions equal to `'alice'` under `NOCASE`, +such as `'alice'` and `'Alice'`, and then those equal to `'bob'`, and returns +up to `k` rows from each, and the query fails if it uses `LIMIT` instead of +`k = ?`. SQLite first keeps only one of the values that the collation treats +as equal, so `name collate nocase in ('alice', 'Alice')` runs one `=`, which +searches both of those partitions. An `IS NOT` condition with `COLLATE` on the value, as in `name is not 'alice' collate nocase`, can return fewer than `k` rows, because `vec0` compares it bytewise and SQLite drops each returned row that equals the value under that collation. @@ -193,9 +203,9 @@ conditions include: `name collate nocase != 'alice'` - `NOT IN`, `NOT LIKE`, and `NOT GLOB` - `LIKE ... ESCAPE`. For some patterns that begin with literal text, such as - `'a!_%' escape '!'`, SQLite also passes `vec0` range conditions, which `vec0` - compares bytewise, so the query can return fewer than `k` rows even with - `LIMIT`, without an error. + `'a!_%' escape '!'`, SQLite instead passes `vec0` range conditions that + select the same rows, so the query returns up to `k` matching rows and works + with `LIMIT`. - `IS TRUE`, `IS FALSE`, `IS NOT TRUE`, and `IS NOT FALSE` - conditions joined with `OR`, except `=` conditions on one column, which SQLite turns into `IN` diff --git a/sqlite-vec.c b/sqlite-vec.c index e7b25513..dd1c150e 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -4066,6 +4066,11 @@ struct vec0_vtab { // Will change the schema of the _rowids table, and insert/query logic. int pkIsText; + // The database's text encoding, SQLITE_UTF8, SQLITE_UTF16LE or + // SQLITE_UTF16BE. vec0 keeps text metadata as UTF-8, and SQLite orders text + // under BINARY by the bytes of the database's encoding. + int textEncoding; + // True if the table declares the hidden command column named after the // table, which runs `INSERT INTO t(t) VALUES ('optimize')`. False when // another column has the table's name; see vec0_init(). @@ -5369,6 +5374,40 @@ void vec0_cursor_clear(vec0_cursor *pCur) { } } +// Reads the database's text encoding into *pEncoding: SQLITE_UTF8, +// SQLITE_UTF16LE or SQLITE_UTF16BE. Every attached database shares the main +// database's encoding. PRAGMA encoding names it "UTF-8", "UTF-16le" or +// "UTF-16be", and a build with SQLITE_OMIT_UTF16 has no such pragma, so it +// returns no row there, where every database is UTF-8 (PragTyp_ENCODING in +// https://sqlite.org/src/file/src/pragma.c and sqlite3InitOne() in +// https://sqlite.org/src/file/src/prepare.c). +static int vec0_read_text_encoding(sqlite3 *db, int *pEncoding) { + sqlite3_stmt *stmt = NULL; + int rc = sqlite3_prepare_v2(db, "PRAGMA encoding", -1, &stmt, NULL); + if (rc != SQLITE_OK) { + return rc; + } + *pEncoding = SQLITE_UTF8; + rc = sqlite3_step(stmt); + if (rc == SQLITE_ROW) { + const char *zEncoding = (const char *)sqlite3_column_text(stmt, 0); + if (!zEncoding) { + rc = SQLITE_NOMEM; + } else { + if (sqlite3_stricmp(zEncoding, "UTF-16le") == 0) { + *pEncoding = SQLITE_UTF16LE; + } else if (sqlite3_stricmp(zEncoding, "UTF-16be") == 0) { + *pEncoding = SQLITE_UTF16BE; + } + rc = SQLITE_OK; + } + } else if (rc == SQLITE_DONE) { + rc = SQLITE_OK; + } + sqlite3_finalize(stmt); + return rc; +} + #define VEC_CONSTRUCTOR_ERROR "vec0 constructor error: " static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv, sqlite3_vtab **ppVtab, char **pzErr, bool isCreate) { @@ -5377,6 +5416,23 @@ static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv, int rc; const char *zSql; + // The text filters need the encoding, so a failed read fails the + // constructor; an interrupted one reports as an interrupt, as + // vec0_interrupt_result() makes the other methods do. + int textEncoding; + rc = vec0_read_text_encoding(db, &textEncoding); + if (rc != SQLITE_OK) { + if (rc == SQLITE_INTERRUPT) { + *pzErr = sqlite3_mprintf("%s", sqlite3_errstr(rc)); + } else { + *pzErr = + sqlite3_mprintf(VEC_CONSTRUCTOR_ERROR + "could not read the database's text encoding: %s", + sqlite3_errmsg(db)); + } + return vec0_error_rc(rc); + } + pNew = sqlite3_malloc(sizeof(*pNew)); if (pNew == 0) return SQLITE_NOMEM; @@ -5725,6 +5781,7 @@ static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv, pNew->db = db; pNew->pkIsText = pkColumnType == SQLITE_TEXT; + pNew->textEncoding = textEncoding; pNew->hasCommandColumn = hasCommandColumn; pNew->schemaName = sqlite3_mprintf("%s", schemaName); if (!pNew->schemaName) { @@ -6442,10 +6499,9 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { int partition_idx = vec0_column_idx_to_partition_idx(p, iColumn); char value = 0; int omit = 1; - // vec0 compares partition keys bytewise. An IN or IS that names another - // collation keeps the handling it had before vec0 consumed them: an IN - // stays an `=` that SQLite runs once per value, and SQLite applies an IS - // to the rows vec0 returns. + // An IN or IS that names a collation other than BINARY keeps the handling + // it had before vec0 consumed them: an IN stays an `=` that SQLite runs + // once per value, and SQLite applies an IS to the rows vec0 returns. int isBinary = sqlite3_stricmp(sqlite3_vtab_collation(pIdxInfo, i), "BINARY") == 0; @@ -6512,6 +6568,43 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { } } + // The _chunks query compares a text key under the collation the query + // names, as the metadata constraints below do. SQLite reports BINARY for + // != whatever collation the query names, so vec0 compares it bytewise + // and SQLite re-checks each row, as for IS NOT. + char collation = VEC0_METADATA_COLLATION_BINARY; + if (value && p->paritition_columns[partition_idx].type == SQLITE_TEXT) { + switch (value) { + case VEC0_PARTITION_OPERATOR_EQ: + case VEC0_PARTITION_OPERATOR_GT: + case VEC0_PARTITION_OPERATOR_LE: + case VEC0_PARTITION_OPERATOR_LT: + case VEC0_PARTITION_OPERATOR_GE: { + const char *zCollation = sqlite3_vtab_collation(pIdxInfo, i); + if (sqlite3_stricmp(zCollation, "NOCASE") == 0) { + collation = VEC0_METADATA_COLLATION_NOCASE; + } else if (sqlite3_stricmp(zCollation, "RTRIM") == 0) { + collation = VEC0_METADATA_COLLATION_RTRIM; + } else if (sqlite3_stricmp(zCollation, "BINARY") != 0) { + rc = SQLITE_ERROR; + vtab_set_error(pVTab, + "Collation \"%s\" is not supported on the \"%.*s\" " + "partition key in a vec0 KNN query. Only BINARY, " + "NOCASE, and RTRIM are supported.", + zCollation, + p->paritition_columns[partition_idx].name_length, + p->paritition_columns[partition_idx].name); + goto done; + } + break; + } + case VEC0_PARTITION_OPERATOR_NE: { + omit = 0; + break; + } + } + } + if (value) { pIdxInfo->aConstraintUsage[i].argvIndex = argvIndex++; pIdxInfo->aConstraintUsage[i].omit = omit; @@ -6519,7 +6612,7 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { VEC0_IDXSTR_KIND_KNN_PARTITON_CONSTRAINT); sqlite3_str_appendchar(idxStr, 1, 'A' + partition_idx); sqlite3_str_appendchar(idxStr, 1, value); - sqlite3_str_appendchar(idxStr, 1, '_'); + sqlite3_str_appendchar(idxStr, 1, collation); } } @@ -7095,7 +7188,7 @@ int vec0_chunks_iter(vec0_vtab *p, const char *idxStr, int argc, int partition_idx = idxStr[idx + 1] - 'A'; int operator = idxStr[idx + 2]; - // idxStr[idx + 3] is just null, a '_' placeholder + char collation = idxStr[idx + 3]; if (!appendedWhere) { sqlite3_str_appendall(s, " WHERE "); @@ -7164,6 +7257,11 @@ int vec0_chunks_iter(vec0_vtab *p, const char *idxStr, int argc, return SQLITE_ERROR; } } + if (collation == VEC0_METADATA_COLLATION_NOCASE) { + sqlite3_str_appendall(s, "COLLATE NOCASE "); + } else if (collation == VEC0_METADATA_COLLATION_RTRIM) { + sqlite3_str_appendall(s, "COLLATE RTRIM "); + } } char *zSql = sqlite3_str_finish(s); @@ -7416,12 +7514,96 @@ static int vec0_text_binary_cmp(const char *a, int na, const char *b, int nb) { return (na > nb) - (na < nb); } -// Compares text the way SQLite's BINARY, NOCASE, or RTRIM collation does. -// NOCASE folds ASCII letters only, and sqlite3_strnicmp() stops at a NUL both -// values have at the same offset. RTRIM ignores trailing spaces, but no other -// trailing whitespace. -static int vec0_text_cmp(vec0_metadata_collation collation, const char *a, - int na, const char *b, int nb) { +// The UTF-8 conversion below is SQLite's own, from READ_UTF8, WRITE_UTF16LE, +// WRITE_UTF16BE and sqlite3Utf8Trans1[] in +// https://sqlite.org/src/file/src/utf.c, so a value converts to the bytes +// SQLite stores for it in a UTF-16 database. +static const unsigned char vec0Utf8Trans1[] = { + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, + 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, + 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x00, + 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, + 0x0c, 0x0d, 0x0e, 0x0f, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, + 0x07, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x00, 0x00, +}; + +// Reads one character of UTF-8, reading an overlong form, a surrogate, U+FFFE +// or U+FFFF as U+FFFD. +static u32 vec0_utf8_read(const unsigned char **pz, const unsigned char *zEnd) { + u32 c = *((*pz)++); + if (c >= 0xc0) { + c = vec0Utf8Trans1[c - 0xc0]; + while (*pz != zEnd && (**pz & 0xc0) == 0x80) { + c = (c << 6) + (0x3f & *((*pz)++)); + } + if (c < 0x80 || (c & 0xFFFFF800) == 0xD800 || (c & 0xFFFFFFFE) == 0xFFFE) { + c = 0xFFFD; + } + } + return c; +} + +// Writes c as UTF-16 in the byte order of encoding, SQLITE_UTF16LE or +// SQLITE_UTF16BE, and returns the number of bytes, 2 or 4. +static int vec0_utf16_write(unsigned char *z, u32 c, int encoding) { + // the offsets of the high and low byte within each 16-bit unit + int hi = encoding == SQLITE_UTF16BE ? 0 : 1; + int lo = 1 - hi; + if (c <= 0xFFFF) { + z[hi] = (unsigned char)((c >> 8) & 0x00FF); + z[lo] = (unsigned char)(c & 0x00FF); + return 2; + } + z[hi] = (unsigned char)(0x00D8 + (((c - 0x10000) >> 18) & 0x03)); + z[lo] = + (unsigned char)(((c >> 10) & 0x003F) + (((c - 0x10000) >> 10) & 0x00C0)); + z[2 + hi] = (unsigned char)(0x00DC + ((c >> 8) & 0x03)); + z[2 + lo] = (unsigned char)(c & 0x00FF); + return 4; +} + +// Compares UTF-8 text the way SQLite's BINARY collation compares it in a +// database whose encoding is SQLITE_UTF16LE or SQLITE_UTF16BE: memcmp() over +// the bytes of each value's UTF-16 form, then the shorter value sorts first. +static int vec0_text_utf16_cmp(int encoding, const char *a, int na, + const char *b, int nb) { + const unsigned char *za = (const unsigned char *)a; + const unsigned char *zb = (const unsigned char *)b; + const unsigned char *zaEnd = za + na; + const unsigned char *zbEnd = zb + nb; + unsigned char ua[4], ub[4]; + int ia = 0, nua = 0, ib = 0, nub = 0; + for (;;) { + if (ia == nua && za != zaEnd) { + nua = vec0_utf16_write(ua, vec0_utf8_read(&za, zaEnd), encoding); + ia = 0; + } + if (ib == nub && zb != zbEnd) { + nub = vec0_utf16_write(ub, vec0_utf8_read(&zb, zbEnd), encoding); + ib = 0; + } + int aDone = ia == nua; + int bDone = ib == nub; + if (aDone || bDone) { + return bDone - aDone; + } + if (ua[ia] != ub[ib]) { + return ua[ia] - ub[ib]; + } + ia++; + ib++; + } +} + +// Compares text the way SQLite's BINARY, NOCASE, or RTRIM collation does in a +// database of the given text encoding. NOCASE folds ASCII letters only, and +// sqlite3_strnicmp() stops at a NUL both values have at the same offset. RTRIM +// ignores trailing spaces, but no other trailing whitespace. SQLite registers +// BINARY in each encoding and NOCASE and RTRIM for UTF-8 only, so those two +// compare UTF-8 in every encoding (openDatabase() in +// https://sqlite.org/src/file/src/main.c). +static int vec0_text_cmp(vec0_metadata_collation collation, int encoding, + const char *a, int na, const char *b, int nb) { switch (collation) { case VEC0_METADATA_COLLATION_NOCASE: { int cmp = sqlite3_strnicmp(a, b, min(na, nb)); @@ -7437,6 +7619,9 @@ static int vec0_text_cmp(vec0_metadata_collation collation, const char *a, nb--; break; case VEC0_METADATA_COLLATION_BINARY: + if (encoding != SQLITE_UTF8) { + return vec0_text_utf16_cmp(encoding, a, na, b, nb); + } break; } return vec0_text_binary_cmp(a, na, b, nb); @@ -7477,8 +7662,8 @@ static int vec0_metadata_filter_text_whole( } } for (size_t j = 0; j < nTargets; j++) { - int cmp = - vec0_text_cmp(collation, s, n, targets[j].zString, targets[j].n); + int cmp = vec0_text_cmp(collation, p->textEncoding, s, n, + targets[j].zString, targets[j].n); if (vec0_metadata_operator_holds(op, cmp)) { bitmap_set(b, i, 1); break; diff --git a/tests/test-metadata.py b/tests/test-metadata.py index fa72a019..8bca9d25 100644 --- a/tests/test-metadata.py +++ b/tests/test-metadata.py @@ -2,6 +2,7 @@ import sqlite3 from collections import OrderedDict import json +from conftest import get_extension_path def test_constructor_limit(db, snapshot): @@ -1615,6 +1616,124 @@ def sql(where, parameters=[]): assert mismatches == [] +@pytest.mark.parametrize("encoding", ["UTF-16le", "UTF-16be"]) +def test_knn_text_filters_in_utf16_database(encoding): + # SQLite orders text under BINARY by the bytes of the database's encoding, + # which orders some characters differently in UTF-16 than in UTF-8, and + # vec0 compares in that order, so a filter keeps the plain table's rows + # before the k nearest are chosen and works with LIMIT. A literal GLOB or + # LIKE pattern adds the range conditions SQLite derives from its prefix. + db = sqlite3.connect(":memory:") + db.enable_load_extension(True) + db.load_extension(get_extension_path()) + db.enable_load_extension(False) + db.execute(f"pragma encoding = '{encoding}'") + db.execute( + "create virtual table v using vec0(vector float[1], t text, chunk_size=8)" + ) + db.execute("create table plain(id integer primary key, t text)") + texts = [ + "a", + "Ā", + "￿", + "\U0001f600", + "z", + "é", + "abcdefghijkl￿", + "abcdefghijkl\U0001f600", + "abc", + "ABC", + "abd", + "abĀ", + ] + # row i's vector is [i], so a query for [0] ranks rows by rowid + for i, t in enumerate(texts, 1): + db.execute("insert into v(rowid, vector, t) values (?, ?, ?)", [i, f"[{i}]", t]) + db.execute("insert into plain values (?, ?)", [i, t]) + + def plain(where, parameters): + return [ + r[0] + for r in db.execute( + f"select id from plain where {where} order by id", parameters + ) + ] + + def knn(where, parameters, form, k): + return [ + r[0] + for r in db.execute( + f"select rowid from v where vector match '[0]' and {where} {form}", + parameters + [k], + ) + ] + + mismatches = [] + for op in ["=", "!=", "<", "<=", ">", ">="]: + for collation in ["", " collate nocase", " collate rtrim"]: + for value in texts: + where = f"t {op} ?{collation}" + want = plain(where, [value]) + got = sorted(knn(where, [value], "and k = ?", len(texts))) + if got != want: + mismatches.append((where, value, got, want)) + if op == "!=" and collation: + # SQLite checks a collated != again, so it can drop rows + continue + for k in range(1, len(texts) + 1): + for form in ["and k = ?", "limit ?"]: + got = knn(where, [value], form, k) + if got != want[:k]: + mismatches.append((where, value, form, k, got)) + literal_patterns = [ + f"t like '{p}'" for p in ["ab%", "AB%", "abc%", "abcdefghijkl%", "Ā%", "é%"] + ] + [f"t glob '{p}'" for p in ["ab*", "abc*", "abcdefghijkl*", "Ā*", "a[b]*"]] + for case_sensitive_like in [0, 1]: + db.execute(f"pragma case_sensitive_like = {case_sensitive_like}") + for where in literal_patterns: + want = plain(where, []) + for k in range(1, len(texts) + 1): + for form in ["and k = ?", "limit ?"]: + got = knn(where, [], form, k) + if got != want[:k]: + mismatches.append((where, case_sensitive_like, form, k, got)) + assert mismatches == [] + + +def test_knn_fails_to_open_when_encoding_is_unreadable(tmp_path): + # vec0 reads the database's text encoding with PRAGMA encoding when it + # opens a table, so a connection whose authorizer denies the pragma fails + # to open it instead of comparing text in a guessed encoding. + path = str(tmp_path / "encoding.db") + db = sqlite3.connect(path) + db.enable_load_extension(True) + db.load_extension(get_extension_path()) + db.execute("create virtual table v using vec0(vector float[1], t text)") + db.execute( + "insert into v(rowid, vector, t) values (1, '[1]', 'a'), (2, '[2]', 'z')" + ) + db.commit() + db.close() + + db = sqlite3.connect(path) + db.enable_load_extension(True) + db.load_extension(get_extension_path()) + db.set_authorizer( + lambda action, *args: ( + sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_PRAGMA + else sqlite3.SQLITE_OK + ) + ) + with pytest.raises( + sqlite3.OperationalError, + match="could not read the database's text encoding: not authorized", + ): + db.execute( + "select rowid from v where vector match '[0]' and k = 1 and t > 'm'" + ).fetchall() + + def test_knn_text_filter_null_long_value(db): # A text value longer than the 12-byte prefix kept in the chunk is compared # in full, which once read through the NULL pointer of a NULL constraint diff --git a/tests/test-partition-keys.py b/tests/test-partition-keys.py index c0acd3d0..fb7e13ee 100644 --- a/tests/test-partition-keys.py +++ b/tests/test-partition-keys.py @@ -1,6 +1,8 @@ import sqlite3 from collections import OrderedDict +import pytest + def test_constructor_limit(db, snapshot): assert ( @@ -216,12 +218,12 @@ def test_knn_partition_filters_match_plain_table(db): def test_knn_partition_filters_with_collation(db): - # vec0 compares partition keys bytewise. SQLite reports a text IS - # constraint's collation, so vec0 leaves one that names another collation - # to SQLite, which applies it to the k rows vec0 returns. It reports BINARY - # for IS NOT whatever collation the query names, so vec0 filters bytewise - # and SQLite re-checks each row. Either way a query returns the nearest of - # the plain table's rows, possibly fewer than k, so this asserts a prefix. + # SQLite reports a text IS constraint's collation, so vec0 leaves one that + # names a collation other than BINARY to SQLite, which applies it to the k + # rows vec0 returns. It reports BINARY for IS NOT whatever collation the + # query names, so vec0 filters bytewise and SQLite re-checks each row. + # Either way a query returns the nearest of the plain table's rows, + # possibly fewer than k, so this asserts a prefix. _create_partition_filter_tables(db) conditions = [] for collation in ["nocase", "rtrim"]: @@ -238,37 +240,87 @@ def test_knn_partition_filters_with_collation(db): assert len(actual) <= k -def test_knn_partition_in_with_collation_keeps_bytewise_matches(db): - # vec0 runs an IN that names a collation other than BINARY as one bytewise - # `=` per value, as it did before it consumed IN. When only one value has - # matches and they have no case or trailing-space variants, that returns - # the plain table's rows; leaving the IN to SQLite would apply it to the k - # nearest "other" rows and return none. +def test_knn_partition_comparisons_with_collation(db): + # The _chunks query compares a text key under the collation the query + # names. SQLite reports BINARY for != whatever collation the query names, + # so vec0 compares it bytewise and SQLite re-checks each row: a query + # returns the nearest of the plain table's rows, possibly fewer than k. + _create_partition_filter_tables(db) + values = ["a", "A", "abc", "ABC", "abc ", "é", "É", "b", 1] + for collation in ["nocase", "rtrim"]: + for op in ["=", "<", "<=", ">", ">="]: + for value in values: + for condition in [ + f"p {op} ? collate {collation}", + f"p collate {collation} {op} ?", + ]: + _assert_knn_partition_filter_matches_plain(db, condition, [value]) + condition = f"p != ? collate {collation}" + for value in values: + expected = _plain_partition_filter(db, condition, [value]) + for k in range(1, len(PARTITION_FILTER_ROWS) + 1): + actual = _knn_partition_filter(db, condition, [value], "k", k) + assert actual == expected[: len(actual)], (condition, value, k) + assert len(actual) <= k + + +def test_knn_partition_like_escape_with_literal_prefix(db): + # For a LIKE ... ESCAPE pattern that is literal text followed by %, SQLite + # passes vec0 range conditions that select the rows the LIKE matches, under + # NOCASE unless case_sensitive_like is on, so the query works with LIMIT. + _create_partition_filter_tables(db) + for case_sensitive_like in [0, 1]: + db.execute(f"pragma case_sensitive_like = {case_sensitive_like}") + for pattern in ["ab%", "AB%", "!_x%", "!%x%", "a%", "é%"]: + _assert_knn_partition_filter_matches_plain( + db, f"p like '{pattern}' escape '!'", [] + ) + + +def test_knn_partition_comparison_with_other_collation_fails(db): + _create_partition_filter_tables(db) + db.create_collation("reverse", lambda a, b: (a < b) - (a > b)) + with pytest.raises(sqlite3.OperationalError, match='Collation "reverse"'): + _knn_partition_filter(db, "p = ? collate reverse", ["a"], "k", 3) + + +def test_knn_partition_in_with_collation_compares_each_value_under_it(db): + # vec0 runs an IN that names a collation other than BINARY as one `=` per + # value under that collation, and SQLite first keeps one of the values the + # collation treats as equal. When one value has matches, that returns the + # plain table's rows, its case variants under NOCASE and its trailing-space + # variants under RTRIM included; leaving the IN to SQLite would apply it to + # the k nearest "other" rows and return none. db.execute("create virtual table v using vec0(p text partition key, e float[1])") db.execute("create table plain(p)") - values = ["other"] * 6 + ["acme"] * 3 + ["beta"] * 3 + values = ["other"] * 6 + ["acme", "ACME", "acme ", "ACME ", "beta", "acme"] for rowid, value in enumerate(values, 1): db.execute( "insert into v(rowid, p, e) values (?, ?, ?)", [rowid, value, f"[{rowid}]"] ) db.execute("insert into plain(rowid, p) values (?, ?)", [rowid, value]) - for collation in ["nocase", "rtrim"]: - condition = f"p collate {collation} in ('acme', 'zzz')" - expected = [ - row[0] - for row in db.execute( - f"select rowid from plain where {condition} order by rowid" - ) - ] - for k in range(1, len(values) + 1): - actual = [ + for collation, lists in [ + ("nocase", ["('acme', 'zzz')", "('ACME', 'zzz')", "('acme', 'ACME')"]), + ("rtrim", ["('acme', 'zzz')", "('acme ', 'zzz')", "('acme', 'acme ')"]), + ]: + for values_list in lists: + condition = f"p collate {collation} in {values_list}" + expected = [ row[0] for row in db.execute( - f"select rowid from v where e match '[0]' and {condition} and k = ?", - [k], + f"select rowid from plain where {condition} order by rowid" ) ] - assert actual == expected[:k], (condition, k) + assert len(expected) == 3, condition + for k in range(1, len(values) + 1): + actual = [ + row[0] + for row in db.execute( + f"select rowid from v where e match '[0]' and {condition} and k = ?", + [k], + ) + ] + assert actual == expected[:k], (condition, k) class Row: