fix(sqlite-vec.c): sum cosine magnitudes in double precision

Squared float32 elements below about 3e-23 or above about 2e19 underflowed or
overflowed float, giving wrong cosine distances and normalizations
(asg017#324, vlasky/sqlite-vec@d75c245).

Author: mceachen

CHANGELOG.md

@@ -18,6 +18,7 @@
 
 ### Fixed
 
+- Fixed cosine distance between `float32` vectors whose elements are very small or very large. Squared magnitudes were summed in `float`, so elements below about 3e-23 made a magnitude of 0, which this fork reports as distance 1.0, and elements above about 2e19 made it Inf, so `vec_distance_cosine()`, cosine KNN queries, and MMR reranking returned wrong distances for those vectors. `vec_normalize()` had the same problem. Both now sum in `double`, so ordinary cosine distances can differ from 2.0.2 in the last few digits. Ported from [vlasky@d75c245](https://github.com/vlasky/sqlite-vec/commit/d75c245) ([asg017#324](https://github.com/asg017/sqlite-vec/issues/324)).
 - Fixed inserts into `vec0` tables created before v0.2.0-alpha after `optimize` had removed a chunk. Those tables declare their `_vector_chunksNN` and `_metadatachunksNN` shadow tables with `rowid PRIMARY KEY` instead of `rowid INTEGER PRIMARY KEY`, so a new chunk's SQLite rowid could differ from its chunk id, and the insert failed with `Error opening vector blob`; KNN and point queries on the table then failed too. New chunks now set both ids. Ported from [vlasky@3a64182](https://github.com/vlasky/sqlite-vec/commit/3a64182).
 - An `INSERT` into a `vec0` table with a value of the wrong type for a metadata or auxiliary column now fails before writing anything. It used to fail after writing the row's rowid, vector, and some column values, and a surrounding transaction could commit that partial row, because `vec0` writes its shadow tables through separate statements that a failed `INSERT` does not roll back.
 - Fixed memory and handle leaks that [vlasky@5ae4fed](https://github.com/vlasky/sqlite-vec/commit/5ae4fed) fixed in its fork and this fork still had: every `UPDATE` of an auxiliary column leaked its SQL string; a `vec0` constructor that failed after parsing some columns leaked their names; a long-text metadata write that failed left its blob handle open, so `sqlite3_close()` returned `SQLITE_BUSY`; a KNN query leaked its top-k buffers when MMR reranking failed; and text-metadata filtering leaked a rowid buffer when its blob read failed.

sqlite-vec.c

@@ -604,21 +604,25 @@ static f32 distance_cosine_float(const void *pVect1v, const void *pVect2v,
   f32 *pVect2 = (f32 *)pVect2v;
   size_t qty = *((size_t *)qty_ptr);
 
-  f32 dot = 0;
-  f32 aMag = 0;
-  f32 bMag = 0;
+  // Accumulate in double: the square of a small (or large) f32 element
+  // underflows (or overflows) f32, which made the magnitude 0 (or Inf) and
+  // the distance +-Inf, NaN or plain wrong for nonzero finite vectors. Every
+  // f32 product is representable in double.
+  double dot = 0;
+  double aMag = 0;
+  double bMag = 0;
   for (size_t i = 0; i < qty; i++) {
-    dot += *pVect1 * *pVect2;
-    aMag += *pVect1 * *pVect1;
-    bMag += *pVect2 * *pVect2;
+    dot += (double)*pVect1 * (double)*pVect2;
+    aMag += (double)*pVect1 * (double)*pVect1;
+    bMag += (double)*pVect2 * (double)*pVect2;
     pVect1++;
     pVect2++;
   }
   // Handle zero vectors: return max distance (1.0) to avoid division by zero
   if (aMag == 0 || bMag == 0) {
     return 1.0f;
   }
-  return 1 - (dot / (sqrtf(aMag) * sqrtf(bMag)));
+  return 1 - (dot / (sqrt(aMag) * sqrt(bMag)));
 }
 static f32 distance_cosine_int8(const void *pA, const void *pB,
                                 const void *pD) {
@@ -1973,9 +1977,11 @@ static void vec_normalize(sqlite3_context *context, int argc,
 
   f32 *v = (f32 *)vector;
 
-  f32 norm = 0;
+  // Accumulate in double so that squaring small or large f32 elements does
+  // not underflow or overflow, see distance_cosine_float().
+  double norm = 0;
   for (size_t i = 0; i < dimensions; i++) {
-    norm += v[i] * v[i];
+    norm += (double)v[i] * (double)v[i];
   }
   norm = sqrt(norm);
   for (size_t i = 0; i < dimensions; i++) {

tests/__snapshots__/test-knn-mmr.ambr

@@ -9,19 +9,19 @@
       }),
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 3,
-        'distance': 0.011511921882629395,
+        'distance': 0.011511985212564468,
       }),
       OrderedDict({
         'rowid': 4,
-        'distance': 0.02060168981552124,
+        'distance': 0.020601652562618256,
       }),
       OrderedDict({
         'rowid': 5,
-        'distance': 0.03227579593658447,
+        'distance': 0.032275766134262085,
       }),
     ]),
   })
@@ -36,19 +36,19 @@
       }),
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 8,
         'distance': 1.0,
       }),
       OrderedDict({
         'rowid': 3,
-        'distance': 0.011511921882629395,
+        'distance': 0.011511985212564468,
       }),
       OrderedDict({
         'rowid': 4,
-        'distance': 0.02060168981552124,
+        'distance': 0.020601652562618256,
       }),
     ]),
   })
@@ -63,11 +63,11 @@
       }),
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 3,
-        'distance': 0.020195960998535156,
+        'distance': 0.020195940509438515,
       }),
     ]),
   })
@@ -82,11 +82,11 @@
       }),
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 3,
-        'distance': 0.020195960998535156,
+        'distance': 0.020195940509438515,
       }),
     ]),
   })
@@ -101,7 +101,7 @@
       }),
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 5,
@@ -209,15 +209,15 @@
     'rows': list([
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 5,
         'distance': 1.0,
       }),
       OrderedDict({
         'rowid': 3,
-        'distance': 0.020195960998535156,
+        'distance': 0.020195940509438515,
       }),
     ]),
   })
@@ -232,7 +232,7 @@
       }),
       OrderedDict({
         'rowid': 2,
-        'distance': 0.005062758922576904,
+        'distance': 0.005062811076641083,
       }),
       OrderedDict({
         'rowid': 5,

tests/test-loadable.py

@@ -421,15 +421,7 @@ def check(a, b, dtype=np.float32):
     check([1.2, 0.1], [0.4, -0.4])
     check([-1.2, -0.1], [-0.4, 0.4])
     check([1, 2, 3], [-9, -8, -7], dtype=np.int8)
-    # Use isclose since sqrtf precision differs from sqrt
-    assert isclose(
-        vec_distance_cosine("[1.1, 1.0]", "[1.2, 1.2]"),
-        npy_cosine(
-            np.array([1.1, 1.0], dtype=np.float32),
-            np.array([1.2, 1.2], dtype=np.float32),
-        ),
-        abs_tol=1e-6,
-    )
+    assert vec_distance_cosine("[1.1, 1.0]", "[1.2, 1.2]") == 0.0011318627512082458
 
     vec_distance_cosine_bit = lambda *args: db.execute(
         "select vec_distance_cosine(vec_bit(?), vec_bit(?))", args
@@ -485,6 +477,29 @@ def test_vec_distance_cosine_zero_vector():
     assert vec_distance_cosine_bit(b"\x00", b"\x00") == 1.0
 
 
+def test_vec_distance_cosine_extreme_magnitudes():
+    # squaring very small or very large f32 elements used to underflow or
+    # overflow the f32 magnitude accumulators, returning +-Inf, NaN or a wrong
+    # distance for nonzero finite vectors
+    # https://github.com/asg017/sqlite-vec/issues/324
+    vec_distance_cosine = lambda a, b: db.execute(
+        "select vec_distance_cosine(?, ?)", [_f32(a), _f32(b)]
+    ).fetchone()[0]
+
+    for scale in [1e-45, 1e-30, 1e-25, 1e-20, 1.0, 1e20, 3e38]:
+        # cosine distance is invariant under positive scaling
+        assert vec_distance_cosine([1, 0], [scale, 0]) == 0.0
+        assert vec_distance_cosine([scale, 0], [1, 0]) == 0.0
+        assert vec_distance_cosine([scale, 0], [scale, 0]) == 0.0
+        assert vec_distance_cosine([1, 0], [-scale, 0]) == 2.0
+        assert vec_distance_cosine([1, 0], [0, scale]) == 1.0
+        assert isclose(
+            vec_distance_cosine([1, 0], [scale, scale]),
+            1 - 1 / np.sqrt(2),
+            abs_tol=1e-6,
+        )
+
+
 def test_ensure_vector_match_cleanup_on_second_vector_error():
     """
     Test that ensure_vector_match properly cleans up the first vector
@@ -756,6 +771,18 @@ def test_vec_normalize():
         -0.6324555277824402,
     ]
 
+    # small and large nonzero vectors are normalized rather than treated as
+    # zero-magnitude or returned as a zero vector, the squared magnitude used
+    # to underflow/overflow f32
+    # https://github.com/asg017/sqlite-vec/issues/324
+    for scale in [1e-45, 1e-30, 1e-25, 1e-22, 1e-20, 1e20, 3e38]:
+        assert vec_normalize(_f32([scale, 0])) == _f32([1, 0])
+        assert vec_normalize(_f32([0, -scale])) == _f32([0, -1])
+        assert np.allclose(
+            np.frombuffer(vec_normalize(_f32([scale, scale])), dtype=np.float32),
+            [1 / np.sqrt(2)] * 2,
+        )
+
 
 def test_vec_slice():
     vec_slice = lambda *args, f="?": db.execute(
@@ -2677,6 +2704,44 @@ def test_vec0_distance_metric():
         assert isclose(actual["distance"], expected["distance"], abs_tol=1e-6)
 
 
+def test_vec0_cosine_extreme_magnitudes():
+    # tiny or huge nonzero vectors used to get +-Inf, NaN or wrong cosine
+    # distances, which ranked them ahead of (or behind) the true nearest
+    # neighbors
+    # https://github.com/asg017/sqlite-vec/issues/324
+    db = connect(EXT_PATH)
+    db.execute("create virtual table v using vec0(a float[2] distance_metric=cosine)")
+    db.executemany(
+        "insert into v(rowid, a) values (?, ?)",
+        [
+            (1, _f32([0.9, 0.1])),
+            (2, _f32([1e-25, 0])),
+            (3, _f32([1e-25, 1e-25])),
+            (4, _f32([-1e-25, 0])),
+            (5, _f32([1e-30, 0])),
+            (6, _f32([3e38, 3e38])),
+        ],
+    )
+
+    expected = [
+        {"rowid": 2, "distance": 0.0},
+        {"rowid": 5, "distance": 0.0},
+        {"rowid": 1, "distance": 0.006116265896707773},
+        {"rowid": 3, "distance": 0.2928932309150696},
+        {"rowid": 6, "distance": 0.2928932309150696},
+        {"rowid": 4, "distance": 2.0},
+    ]
+    for query in ([1, 0], [1e-25, 0], [3e38, 0]):
+        rows = execute_all(
+            db,
+            "select rowid, distance from v where a match ? and k = 6",
+            [_f32(query)],
+        )
+        # the order of equidistant rows is unspecified
+        assert sorted(rows, key=lambda r: (r["distance"], r["rowid"])) == expected
+        assert [r["distance"] for r in rows] == [r["distance"] for r in expected]
+
+
 def test_vec0_vacuum():
     db = connect(EXT_PATH)
     db.execute("create virtual table vec_t using vec0(a float[1]);")