fix(sqlite-vec.c): rank a NaN distance last across chunks too

#2 ranks a NaN distance after every number within a chunk, but
merge_sorted_lists() put a later chunk's NaN ahead of the numbers already
merged, which a vec0 table written before 2.1.0 can still produce. Its
CHANGELOG entry now also says what changes on such tables.

Author: mceachen

CHANGELOG.md

@@ -7,7 +7,7 @@
 - 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.
-- A KNN query picks each chunk's k nearest rows with the heap of upstream's `SQLITE_VEC_EXPERIMENTAL_MIN_IDX` build option ([asg017/sqlite-vec@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f4570c1d65a0699e8e38d6c5562da9e23e)) instead of one pass over the chunk for every row it keeps, ordered so that it returns the same rows in the same order as before: on a 10,000-row `float[768]` table, `k = 1000` and `k = 4096` take 9.3 ms instead of 17.8 and 18.0 ms, `k = 100` 8.8 ms instead of 9.2, and `k = 1` 8.7 ms either way. A NaN distance, which only a static blob can produce since `vec0` rejects NaN elements, ranks after every number instead of where its position put it.
+- A KNN query picks each chunk's k nearest rows with the heap of upstream's `SQLITE_VEC_EXPERIMENTAL_MIN_IDX` build option ([asg017/sqlite-vec@0de765f](https://github.com/asg017/sqlite-vec/commit/0de765f4570c1d65a0699e8e38d6c5562da9e23e)) instead of one pass over the chunk for every row it keeps, ordered so that it returns the same rows in the same order as before: on a 10,000-row `float[768]` table, `k = 1000` and `k = 4096` take 9.3 ms instead of 17.8 and 18.0 ms, `k = 100` 8.8 ms instead of 9.2, and `k = 1` 8.7 ms either way. A NaN distance ranks after every number instead of where its position put it. A static blob can produce one, and so can a `vec0` table written before 2.1.0, which started rejecting NaN elements. On such a table a NaN row used to push nearer rows out of the result or behind farther ones; queries now return the nearest rows in order, which can leave out a NaN row that a query with `k` smaller than the row count used to return.
 
 ### Fixed
 

sqlite-vec.c

@@ -6840,7 +6840,8 @@ void merge_sorted_lists(f32 *a, i64 *a_rowids, i64 a_length, f32 *b,
       out_rowids[i] = a_rowids[ptrA];
       ptrA++;
     } else {
-      if (a[ptrA] <= b[b_top_idxs[ptrB]]) {
+      // A NaN ranks after every number, as in min_idx()
+      if (a[ptrA] <= b[b_top_idxs[ptrB]] || isnan(b[b_top_idxs[ptrB]])) {
         out[i] = a[ptrA];
         out_rowids[i] = a_rowids[ptrA];
         ptrA++;

tests/test-knn-topk.py

@@ -64,8 +64,8 @@ def knn(k):
 
 
 def test_knn_ranks_a_nan_distance_last(db):
-    # vec0 refuses NaN elements, but a static blob reads its vectors from the
-    # application's memory as they are, so a distance to one can be NaN.
+    # A static blob reads its vectors from the application's memory as they
+    # are, so a distance to one can be NaN.
     db.enable_load_extension(True)
     db.load_extension(get_extension_path(), entrypoint="sqlite3_vec_static_blobs_init")
     vectors = np.array([[np.nan, 0], [1, 0], [2, 0], [3, 0]], dtype=np.float32)
@@ -80,3 +80,39 @@ def test_knn_ranks_a_nan_distance_last(db):
         [_f32([0, 0])],
     )
     assert [r[0] for r in rows] == [1, 2, 3, 0]
+
+
+def test_knn_ranks_a_stored_nan_distance_last_across_chunks(db):
+    # vec0 has rejected NaN elements since 2.1.0, but a table written by an
+    # earlier release can still hold one. Patching a NaN into the stored chunk
+    # stands in for such a table.
+    db.execute("create virtual table v using vec0(e float[1], chunk_size=8)")
+    for i in range(1, 9):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([10 + i])])
+    for i in range(9, 17):
+        db.execute("insert into v(rowid, e) values (?, ?)", [i, _f32([1])])
+    chunk_id, offset = db.execute(
+        "select chunk_id, chunk_offset from v_rowids where rowid = 16"
+    ).fetchone()
+    vectors = bytearray(
+        db.execute(
+            "select vectors from v_vector_chunks00 where rowid = ?", [chunk_id]
+        ).fetchone()[0]
+    )
+    vectors[offset * 4 : offset * 4 + 4] = _f32([float("nan")])
+    db.execute(
+        "update v_vector_chunks00 set vectors = ? where rowid = ?",
+        [bytes(vectors), chunk_id],
+    )
+
+    def knn(k):
+        return [
+            r[0]
+            for r in db.execute(
+                "select rowid from v where e match ? and k = ?", [_f32([0]), k]
+            )
+        ]
+
+    # The NaN, in the second chunk, ranks after the first chunk's rows too.
+    assert knn(16) == [15, 14, 13, 12, 11, 10, 9, 1, 2, 3, 4, 5, 6, 7, 8, 16]
+    assert knn(9) == [15, 14, 13, 12, 11, 10, 9, 1, 2]