mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
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:
committed by
meta-codesync[bot]
parent
613e0acac5
commit
561e6ecdf3
@@ -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
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user