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 https://github.com/facebookresearch/faiss/issues/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: https://github.com/facebookresearch/faiss/pull/5509

Reviewed By: juancarpio27

Differential Revision: D117536893

Pulled By: mnorris11

fbshipit-source-id: 8568437eb596a202262fc455748063fafca202c1
This commit is contained in:
William
2026-08-26 12:48:20 -07:00
committed by meta-codesync[bot]
parent 613e0acac5
commit 561e6ecdf3
2 changed files with 73 additions and 23 deletions
+3 -3
View File
@@ -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,
+70 -20
View File
@@ -74,6 +74,26 @@ class IDRemap(unittest.TestCase):
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 @@ class Shards(unittest.TestCase):
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 @@ class Shards(unittest.TestCase):
# 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")