Fix add_sa_codes silently accepting non-int64 ids, corrupting stored labels (#5171)

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

**Problem 1 — add_sa_codes silently stores wrong ids when passed non-int64 input**

replacement_add_sa_codes in class_wrappers.py passed the ids array directly to swig_ptr without dtype coercion. The C++ signature is add_sa_codes(idx_t n, const uint8_t* codes, const idx_t* xids) where idx_t = int64_t (8 bytes per element). When the caller passed int32 ids — the default output of np.arange(n) on most platforms — swig_ptr returned a pointer to 4-byte data. The C++ layer then read every 8 consecutive bytes as one int64, producing garbage labels for every stored vector. No exception was raised, no warning was emitted.

This is inconsistent with every other ids-accepting wrapper in the same file. replacement_add_with_ids, replacement_remove_ids, and replacement_permute_entries all call np.ascontiguousarray(ids, dtype='int64') before swig_ptr. add_sa_codes was the only exception. Callers using np.arange(n) without an explicit dtype — common in tutorials and quick-start code — silently ingested a corrupt index.

Fix: add ids = np.ascontiguousarray(ids, dtype='int64') before swig_ptr(ids), matching the pattern used by all other ids-accepting wrappers. The call is a no-op for callers already passing int64, so there is no behavior change for correct inputs.

**Problem 2 — sa_encode calls sa_code_size() twice (one redundant SWIG dispatch)**

replacement_sa_encode called self.sa_code_size() once in the if codes is None branch and once in the else branch. Since sa_code_size() is a virtual C++ method accessed via SWIG, this is one avoidable cross-boundary call per sa_encode invocation. The value is a pure function of index state and cannot change within a call.

replacement_sa_decode in the same block already uses the correct pattern, reusing the code size from codes.shape. The inconsistency between sa_encode and sa_decode confirms this was an oversight.

Fix: hoist self.sa_code_size() into code_size = self.sa_code_size() once before the branch, matching the pattern in sa_decode.

Reviewed By: limqiying

Differential Revision: D103384269

fbshipit-source-id: 3f28a18b7e951c24d23db676be475bb6b298df2f

faiss/python/class_wrappers.py

@@ -842,11 +842,12 @@ def handle_Index(the_class):
         n, d = x.shape
         assert d == self.d
         x = np.ascontiguousarray(x, dtype='float32')
+        code_size = self.sa_code_size()
 
         if codes is None:
-            codes = np.empty((n, self.sa_code_size()), dtype=np.uint8)
+            codes = np.empty((n, code_size), dtype=np.uint8)
         else:
-            assert codes.shape == (n, self.sa_code_size())
+            assert codes.shape == (n, code_size)
 
         self.sa_encode_c(n, swig_ptr(x), swig_ptr(codes))
         return codes
@@ -869,10 +870,12 @@ def handle_Index(the_class):
         assert cs == self.sa_code_size()
         codes = _check_dtype_uint8(codes)
 
+        ids_ptr = None
         if ids is not None:
             assert ids.shape == (n,)
-            ids = swig_ptr(ids)
-        self.add_sa_codes_c(n, swig_ptr(codes), ids)
+            ids = np.ascontiguousarray(ids, dtype="int64")
+            ids_ptr = swig_ptr(ids)
+        self.add_sa_codes_c(n, swig_ptr(codes), ids_ptr)
 
     def replacement_permute_entries(self, perm):
         n, = perm.shape

tests/test_swig_wrapper.py

@@ -185,6 +185,62 @@ class TestException(unittest.TestCase):
         #    assert 'could not parse' in str(e)
 
 
+class TestAddSACodes(unittest.TestCase):
+    """Regression tests for add_sa_codes and sa_encode wrapper correctness."""
+
+    D = 32
+    NB = 200
+
+    def setUp(self):
+        rng = np.random.default_rng(42)
+        self.xb = rng.random((self.NB, self.D), dtype=np.float32)
+        self.xtrain = rng.random((500, self.D), dtype=np.float32)
+        # IDMap2 wrapping PQ8x2: outer index stores (id → code) pairs;
+        # inner PQ8x2 encodes vectors via sa_encode / add_sa_codes.
+        self.index = faiss.index_factory(self.D, "IDMap2,PQ8x2")
+        self.index.train(self.xtrain)
+
+    def test_add_sa_codes_int32_ids_stored_correctly(self):
+        """int32 ids must be coerced to int64 before the C++ boundary.
+
+        Without coercion swig_ptr hands the C++ side a pointer to 4-byte
+        int32 data; C++ interprets every 8 bytes as one idx_t, producing
+        garbage ids with no exception."""
+        codes = self.index.index.sa_encode(self.xb)
+
+        ids_int64 = np.arange(self.NB, dtype="int64") + 1000
+        ids_int32 = ids_int64.astype("int32")
+
+        self.index.add_sa_codes(codes, ids_int32)
+
+        stored = faiss.vector_to_array(self.index.id_map)
+        np.testing.assert_array_equal(stored, ids_int64)
+
+    def test_add_sa_codes_int64_ids_unchanged(self):
+        """int64 ids must pass through as-is (np.ascontiguousarray is a no-op
+        when the array is already int64 and contiguous)."""
+        codes = self.index.index.sa_encode(self.xb)
+
+        ids = np.arange(self.NB, dtype="int64") + 500
+        self.index.add_sa_codes(codes, ids)
+
+        stored = faiss.vector_to_array(self.index.id_map)
+        np.testing.assert_array_equal(stored, ids)
+
+    def test_sa_encode_preallocated_matches_auto_allocated(self):
+        """sa_encode with a pre-allocated buffer must return the same codes
+        as the auto-allocated path, exercising both branches of the hoisted
+        sa_code_size() call."""
+        inner = self.index.index
+        codes_auto = inner.sa_encode(self.xb)
+
+        buf = np.empty_like(codes_auto)
+        codes_pre = inner.sa_encode(self.xb, buf)
+
+        np.testing.assert_array_equal(codes_auto, codes_pre)
+        self.assertIs(codes_pre, buf)
+
+
 @unittest.skipIf(faiss.swig_version() < 0x040000, "swig < 4 does not support Doxygen comments")
 class TestDoxygen(unittest.TestCase):