Merge shard results by metric type rather than assuming similarity (#5509)

Summary:
`IndexShardsTemplate::search` picked the comparator for merging per-shard results by testing for `METRIC_L2`:

```cpp
if (this->metric_type == METRIC_L2) {
    merge_knn_results<idx_t, CMin<distance_t, int>>(...);   // smaller is better
} else {
    merge_knn_results<idx_t, CMax<distance_t, int>>(...);   // larger is better
}
```

Everything that is not `METRIC_L2` was therefore merged as a similarity. `METRIC_L1`, `METRIC_Linf` and `METRIC_Lp` return distances, so once results crossed a shard boundary they came back farthest-first, while the same vectors in a single unsharded index came back nearest-first.

This PR selects the comparator with the existing `is_similarity_metric()` helper from `MetricType.h`, which is exactly the distinction the branch needs.

Fixes #5503

## Behaviour change

`is_similarity_metric()` is true only for `METRIC_INNER_PRODUCT` and `METRIC_Jaccard`, so:

| metric | before | after |
| --- | --- | --- |
| `METRIC_L2` | `CMin` | `CMin` (unchanged) |
| `METRIC_INNER_PRODUCT` | `CMax` | `CMax` (unchanged) |
| `METRIC_L1`, `METRIC_Linf`, `METRIC_Lp`, `METRIC_Canberra`, `METRIC_BrayCurtis`, `METRIC_JensenShannon` | `CMax` | **`CMin`** |
| `METRIC_Jaccard` | `CMax` | `CMax` (unchanged) |

`IndexShardsTemplate<IndexBinary>` is also unaffected: `IndexBinary::metric_type` defaults to `METRIC_L2`, which keeps `CMin` as before.

## Reproduction

Against the released `faiss-cpu` 1.15.0 wheel, using the reporter's example:

```python
import faiss, numpy as np

xb0 = np.array([[0., 0.]], dtype='float32')
xb1 = np.array([[10., 0.]], dtype='float32')
xq  = np.array([[1., 0.]], dtype='float32')

flat = faiss.IndexFlat(2, faiss.METRIC_L1)
flat.add(np.vstack([xb0, xb1]))
print(flat.search(xq, 2))

s0 = faiss.IndexFlat(2, faiss.METRIC_L1); s0.add(xb0)
s1 = faiss.IndexFlat(2, faiss.METRIC_L1); s1.add(xb1)
shards = faiss.IndexShards(2, False, True)
shards.add_shard(s0); shards.add_shard(s1)
print(shards.search(xq, 2))
```

```
single IndexFlat  D,I = [[1. 9.]] [[0 1]]     # nearest first, correct
IndexShards       D,I = [[9. 1.]] [[1 0]]     # farthest first, reversed
```

## Test

Adds `Shards::test_shards_distance_metric_ordering` to `tests/test_meta_index.py`. It splits the dataset across three `METRIC_L1` shards and checks that each result row is ordered nearest-first and matches the distances an unsharded `IndexFlat(METRIC_L1)` returns. Distances rather than labels are compared so that equidistant neighbours may be returned in either order.

The test fails on `main` (the rows come back reversed) and passes with this change.

## Testing notes

I reproduced the bug against the released 1.15.0 wheel as shown above, but I was not able to build faiss from source on this machine (Windows, no local C++ toolchain), so I have not executed the C++ build or run the test suite locally. The added test's pass/fail claim above follows from the comparator change rather than from a local run. Please treat CI as the gate, and I am happy to adjust if anything in the suite disagrees.

Pull Request resolved: #5509

Reviewed By: juancarpio27

Differential Revision: D117536893

Pulled By: mnorris11

fbshipit-source-id: 8568437eb596a202262fc455748063fafca202c1

Author: wilyan09007

faiss/IndexShards.cpp

@@ -243,8 +243,8 @@ void IndexShardsTemplate<IndexT>::search(
 
     this->runOnIndex(fn);
 
-    if (this->metric_type == METRIC_L2) {
-        merge_knn_results<idx_t, CMin<distance_t, int>>(
+    if (is_similarity_metric(this->metric_type)) {
+        merge_knn_results<idx_t, CMax<distance_t, int>>(
                 n,
                 k,
                 nshard,
@@ -253,7 +253,7 @@ void IndexShardsTemplate<IndexT>::search(
                 distances,
                 labels);
     } else {
-        merge_knn_results<idx_t, CMax<distance_t, int>>(
+        merge_knn_results<idx_t, CMin<distance_t, int>>(
                 n,
                 k,
                 nshard,

tests/test_meta_index.py

@@ -74,6 +74,26 @@ def test_id_remap_ivf(self):
 
 class Shards(unittest.TestCase):
 
+    def add_flat_shards(
+        self,
+        shard_index,
+        xbase,
+        metric=faiss.METRIC_L2,
+        metric_arg=0,
+        ni=3,
+        wrap=None,
+    ):
+        """Split xbase over ni IndexFlat shards of shard_index. xbase may be
+        None to leave the shards empty, e.g. to fill them via shard_index.add()
+        """
+        for i in range(ni):
+            n = 0 if xbase is None else len(xbase)
+            shard = faiss.IndexFlat(shard_index.d, metric)
+            shard.metric_arg = metric_arg
+            if xbase is not None:
+                shard.add(xbase[i * n // ni : (i + 1) * n // ni])
+            shard_index.add_shard(shard if wrap is None else wrap(shard))
+
     @unittest.skipIf(
         os.name == "posix" and os.uname().sysname == "Darwin",
         "There is a bug in the OpenMP implementation on OSX.",
@@ -92,26 +112,10 @@ def test_shards(self):
         shard_index_threaded = faiss.IndexShards(d, True)  # explicitly threaded
         shard_index_2 = faiss.IndexShards(d, True, False)
 
-        ni = 3
-        # Populate both indexes with the same data
-        for i in range(ni):
-            i0 = int(i * nb / ni)
-            i1 = int((i + 1) * nb / ni)
-
-            # Add to non-threaded index
-            index_nt = faiss.IndexFlatL2(d)
-            index_nt.add(xb[i0:i1])
-            shard_index_nonthreaded.add_shard(index_nt)
-
-            # Add to threaded index
-            index_t = faiss.IndexFlatL2(d)
-            index_t.add(xb[i0:i1])
-            shard_index_threaded.add_shard(index_t)
-
-            # Add to shard_index_2 for the original test logic
-            index_2 = faiss.IndexFlatL2(d)
-            irm = faiss.IndexIDMap(index_2)
-            shard_index_2.add_shard(irm)
+        self.add_flat_shards(shard_index_nonthreaded, xb)
+        self.add_flat_shards(shard_index_threaded, xb)
+        # populated below by the parallel add rather than shard by shard
+        self.add_flat_shards(shard_index_2, None, wrap=faiss.IndexIDMap)
 
         # test parallel add
         shard_index_2.verbose = True
@@ -145,6 +149,52 @@ def test_shards(self):
             # thousands of the nq*k cells, far above this floor.
             assert ndiff < nq * k / 100.0, f"too many mismatches: {ndiff}"
 
+    def test_shards_metrics(self):
+        # IndexShards merges the per-shard results itself, so it has to know
+        # whether the metric ranks by similarity (largest first) or by distance
+        # (smallest first). Only METRIC_L2 used to be treated as a distance, so
+        # every other dis-similarity metric returned the FARTHEST vectors once
+        # results crossed a shard boundary.
+        rs = np.random.RandomState(123)
+        dim, n, nquery, k = 16, 600, 50, 10
+        # Components in [0, 1) keep every metric well-defined: Jaccard needs
+        # positive components and GOWER numeric dimensions in [0, 1].
+        base = rs.rand(n, dim).astype("float32")
+        queries = rs.rand(nquery, dim).astype("float32")
+        p = 1.5  # metric_arg for METRIC_Lp
+
+        for metric in (
+            faiss.METRIC_INNER_PRODUCT,
+            faiss.METRIC_L2,
+            faiss.METRIC_L1,
+            faiss.METRIC_Linf,
+            faiss.METRIC_Lp,
+            faiss.METRIC_Canberra,
+            faiss.METRIC_BrayCurtis,
+            faiss.METRIC_JensenShannon,
+            faiss.METRIC_Jaccard,
+            faiss.METRIC_NaNEuclidean,
+            faiss.METRIC_GOWER,
+        ):
+            with self.subTest(metric=metric):
+                ref_index = faiss.IndexFlat(dim, metric)
+                ref_index.metric_arg = p
+                ref_index.add(base)
+                Dref, _Iref = ref_index.search(queries, k)
+
+                shard_index = faiss.IndexShards(dim, False, True)
+                self.add_flat_shards(shard_index, base, metric, metric_arg=p)
+                D, _I = shard_index.search(queries, k)
+
+                if faiss.is_similarity_metric(metric):
+                    self.assertTrue(np.all(D[:, :-1] >= D[:, 1:]))
+                else:
+                    self.assertTrue(np.all(D[:, :-1] <= D[:, 1:]))
+                # Same neighbors as the unsharded index. Distances are compared
+                # rather than labels so equidistant neighbors may come in
+                # either order.
+                np.testing.assert_array_almost_equal(D, Dref, decimal=5)
+
     def test_shards_ivf(self):
         ds = SyntheticDataset(32, 1000, 100, 20)
         ref_index = faiss.index_factory(ds.d, "IVF32,SQ8")