mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
fix(python): int64 coercions for MapLong2Long + InvertedLists DOWNCAST chain (#5257)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5257 Two Python binding correctness fixes. Both are silent data-corruption bugs that affect callers using standard numpy idioms. ## Fix 1: MapLong2Long missing int64 dtype coercions `handle_MapLong2Long` in `class_wrappers.py` passed `keys` and `vals` directly to `swig_ptr` without int64 coercion. `swig_ptr` returns the raw memory pointer of the numpy array, ignoring its dtype. The C++ `MapLong2Long::add` and `search_multiple` methods take `const int64_t*`, so when callers pass `int32` arrays (e.g., `np.arange(n)`, which defaults to int32 on most platforms), C++ reads every 8 bytes as one `int64_t`. For a 5-element int32 array, this reads 2 int64 values instead of 5, and each int64 is a bit-merge of two adjacent int32 values. The map is silently populated with completely wrong keys and values, with no exception or warning. This is the identical bug class as the `add_sa_codes` int64 coercion fix (D103384269, 2026-05-01), and the identical pattern omission. Every other ids-accepting wrapper in `class_wrappers.py` (`replacement_add_with_ids`, `replacement_remove_ids`, `replacement_permute_entries`) applies `np.ascontiguousarray(..., dtype='int64')` before `swig_ptr`. `handle_MapLong2Long` was added separately and missed this step. **Fix:** Add `np.ascontiguousarray(keys, dtype='int64')` and `np.ascontiguousarray(vals, dtype='int64')` in `replacement_map_add`, and `np.ascontiguousarray(keys, dtype='int64')` in `replacement_map_search_multiple`. ## Fix 2: InvertedLists* DOWNCAST chain missing ArrayInvertedListsPanorama and SliceInvertedLists The `%typemap(out) faiss::InvertedLists *` in `swigfaiss.swig` runs whenever a C++ function returns `faiss::InvertedLists*` (including `downcast_InvertedLists`). Before this fix, `ArrayInvertedListsPanorama` was absent from the DOWNCAST chain, so `dynamic_cast<faiss::ArrayInvertedLists*>` succeeded first (since `ArrayInvertedListsPanorama` IS-A `ArrayInvertedLists`) and the object was returned as `faiss.ArrayInvertedLists`, hiding all Panorama-specific fields (`cum_sums`, `n_levels`, `level_width`, `get_cum_sums()`). `SliceInvertedLists` fell through all entries to the base `DOWNCAST(InvertedLists)`, returning the base type. `ArrayInvertedListsPanorama` and `SliceInvertedLists` are both defined in `InvertedLists.h` (already `%include`d at line 551) and are fully exposed as Python types. The DOWNCAST entries were simply missing. **Fix:** Add `DOWNCAST(ArrayInvertedListsPanorama)` immediately before `DOWNCAST(ArrayInvertedLists)` (subclass must precede parent); add `DOWNCAST(SliceInvertedLists)` before the base `DOWNCAST(InvertedLists)` catch-all. `downcast_InvertedLists` now returns the correct subtype for both. Note: `%typemap(out)` applies to function return types, not raw member field access. Tests use `faiss.downcast_InvertedLists(index.invlists)`, not `index.invlists` directly — the same pattern as `faiss.downcast_index()` for the `Index*` chain. ## Tests 4 new tests in `TestMapLong2Long`: verify int32 keys/vals are stored and looked up correctly, verify int64 is unchanged, verify search_multiple with int32 keys, verify missing key returns -1. 2 new tests in `TestInvertedListsDowncast`: verify `downcast_InvertedLists` returns `ArrayInvertedListsPanorama` for a Panorama index (with Panorama-specific field access), and `SliceInvertedLists` after `replace_invlists` (with `i0`/`i1` field access). Reviewed By: junjieqi Differential Revision: D106074206 fbshipit-source-id: d0f59ac1e27c37add9c21927ce1dc9af9771d17b
This commit is contained in:
committed by
meta-codesync[bot]
parent
f29d8621fa
commit
34eb98936f
@@ -1275,10 +1275,13 @@ def handle_MapLong2Long(the_class):
|
||||
def replacement_map_add(self, keys, vals):
|
||||
n, = keys.shape
|
||||
assert (n,) == vals.shape
|
||||
keys = np.ascontiguousarray(keys, dtype='int64')
|
||||
vals = np.ascontiguousarray(vals, dtype='int64')
|
||||
self.add_c(n, swig_ptr(keys), swig_ptr(vals))
|
||||
|
||||
def replacement_map_search_multiple(self, keys):
|
||||
n, = keys.shape
|
||||
keys = np.ascontiguousarray(keys, dtype='int64')
|
||||
vals = np.empty(n, dtype='int64')
|
||||
self.search_multiple_c(n, swig_ptr(keys), swig_ptr(vals))
|
||||
return vals
|
||||
|
||||
@@ -947,6 +947,7 @@ typedef int cudaDataType_t;
|
||||
}
|
||||
|
||||
%typemap(out) faiss::InvertedLists * {
|
||||
DOWNCAST (ArrayInvertedListsPanorama) // subclass before parent
|
||||
DOWNCAST (ArrayInvertedLists)
|
||||
DOWNCAST (BlockInvertedLists)
|
||||
#ifndef SWIGWIN
|
||||
@@ -954,6 +955,7 @@ typedef int cudaDataType_t;
|
||||
#endif // !SWIGWIN
|
||||
DOWNCAST (VStackInvertedLists)
|
||||
DOWNCAST (HStackInvertedLists)
|
||||
DOWNCAST (SliceInvertedLists)
|
||||
DOWNCAST (MaskedInvertedLists)
|
||||
DOWNCAST (StopWordsInvertedLists)
|
||||
DOWNCAST (CappedInvertedLists)
|
||||
|
||||
@@ -286,3 +286,49 @@ class TestDoxygen(unittest.TestCase):
|
||||
|
||||
self.assertTrue("a template structure for a set of [min|max]-heaps"
|
||||
in maxheap_array.__doc__)
|
||||
|
||||
|
||||
class TestMapLong2Long(unittest.TestCase):
|
||||
def test_add_coerces_keys_and_vals(self):
|
||||
for dtype in (np.int32, np.int64):
|
||||
with self.subTest(dtype=dtype):
|
||||
m = faiss.MapLong2Long()
|
||||
keys = np.array([0, 1, 2, 3, 4], dtype=dtype)
|
||||
vals = np.array([100, 101, 102, 103, 104], dtype=dtype)
|
||||
m.add(keys, vals)
|
||||
for k, v in zip(range(5), range(100, 105)):
|
||||
self.assertEqual(m.search(k), v)
|
||||
|
||||
def test_search_multiple_coerces_int32_keys(self):
|
||||
m = faiss.MapLong2Long()
|
||||
m.add(
|
||||
np.arange(5, dtype=np.int64),
|
||||
np.arange(5, dtype=np.int64) + 300,
|
||||
)
|
||||
result = m.search_multiple(np.array([0, 2, 4], dtype=np.int32))
|
||||
np.testing.assert_array_equal(result, [300, 302, 304])
|
||||
|
||||
|
||||
class TestInvertedListsDowncast(unittest.TestCase):
|
||||
def test_downcast_ArrayInvertedListsPanorama(self):
|
||||
# downcast_InvertedLists() triggers the typemap; raw .invlists access
|
||||
# does not (typemap applies to function returns, not member getters).
|
||||
d, nlist, n_levels = 32, 8, 2
|
||||
index = faiss.IndexIVFFlatPanorama(
|
||||
faiss.IndexFlatL2(d), d, nlist, n_levels
|
||||
)
|
||||
il = faiss.downcast_InvertedLists(index.invlists)
|
||||
self.assertIsInstance(il, faiss.ArrayInvertedListsPanorama)
|
||||
self.assertEqual(il.n_levels, n_levels)
|
||||
|
||||
def test_downcast_SliceInvertedLists(self):
|
||||
d, nlist = 4, 4
|
||||
# code_size = d * sizeof(float)
|
||||
backing = faiss.ArrayInvertedLists(nlist, d * 4)
|
||||
sil = faiss.SliceInvertedLists(backing, 0, nlist)
|
||||
index = faiss.IndexIVFFlat(faiss.IndexFlatL2(d), d, nlist)
|
||||
# own=False: index doesn't delete sil; Python locals keep both alive.
|
||||
index.replace_invlists(sil, False)
|
||||
inner = faiss.downcast_InvertedLists(index.invlists)
|
||||
self.assertIsInstance(inner, faiss.SliceInvertedLists)
|
||||
self.assertEqual((inner.i0, inner.i1), (0, nlist))
|
||||
|
||||
Reference in New Issue
Block a user