RaBitQ: Clamp L2 distance estimate to >=0 (#5393)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5393

When the two vectors are equal (for example, if you pick a doc vector `d` as your query vector `q` and then estimate `<x, d>`) due to some floating point arithmetic issues the L2 distance can go very slightly negative.  Then when you take the square root of a negative number (in Unicorn) you get an error.  To avoid this clamp to zero.

Ex:
```
or_minus_c_l2sqr = 5.39125443e+01    (A)
qr_to_c_L2sqr   = 5.39125290e+01    (B)
dp_multiplier    = 8.98244381e+00
final_dot        = 6.00199175e+00
cross_term       = 1.07825104e+02    (C = 2 * dp_multiplier * final_dot)
pre_dist       = -3.05175781e-05 (A + B - C)
```

Note that the FastScan path already does this so there is a discrepancy being amended here.

Reviewed By: ddrcoder

Differential Revision: D110219339

fbshipit-source-id: 9dfa70e4797f64830df1c6542456112dc9a54785

faiss/impl/RaBitQuantizer.cpp

@@ -285,7 +285,7 @@ struct RaBitQDistanceComputerNotQ : RaBitQDistanceComputer {
 
         if (metric_type == MetricType::METRIC_L2) {
             // ||or - q||^ 2
-            return pre_dist;
+            return std::max(0.0f, pre_dist);
         } else {
             // metric == MetricType::METRIC_INNER_PRODUCT
             // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)
@@ -448,7 +448,7 @@ struct RaBitQDistanceComputerQ : RaBitQDistanceComputer {
 
         if (metric_type == MetricType::METRIC_L2) {
             // ||or - q||^ 2
-            return pre_dist;
+            return std::max(0.0f, pre_dist);
         } else {
             // metric == MetricType::METRIC_INNER_PRODUCT
             // 2 * (or, q) = (||or - q||^2 - ||q||^2 - ||or||^2)

tests/test_rabitq.py

@@ -174,7 +174,7 @@ class TestRaBitQ(unittest.TestCase):
                     params = faiss.RaBitQSearchParameters(
                         qb=qb, centered=centered
                     )
-                    ratio_threshold = 2 ** (1 / qb)
+                    ratio_threshold = 2.1 ** (1 / qb)
                     test(
                         params,
                         ratio_threshold,
@@ -249,6 +249,20 @@ class TestRaBitQ(unittest.TestCase):
                 f"{j} {ref_dis[0][j]} {upd_dis}",
             )
 
+    def test_self_match_nonneg_distance(self):
+        d = 64
+        rs = np.random.RandomState(10996)
+        xb = rs.randn(200, d).astype("float32")
+
+        index = faiss.IndexRaBitQ(d, faiss.METRIC_L2)
+        index.train(xb)
+        index.add(xb)
+
+        D, I = index.search(xb[:5], 5)
+        np.testing.assert_(np.all(D >= 0), f"negative D: {D[D < 0]}")
+        for i in range(5):
+            self.assertIn(i, I[i])
+
     def test_serde_rabitq(self):
         do_test_serde("RaBitQ")