Support IndexIDMap/IndexIDMap2 in reverse_index_factory (#5266)

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

`reverse_index_factory` in `contrib/factory_tools.py` raised `NotImplementedError` for `IndexIDMap` and `IndexIDMap2`, even though both simply wrap an inner index. This adds handling for both: an `IndexIDMap` is described as `IDMap,<inner>` and an `IndexIDMap2` as `IDMap2,<inner>`, recursing into the wrapped index. `IndexIDMap2` is checked first because it subclasses `IndexIDMap`. This matches the C++ `reverse_index_factory` in `faiss/factory_tools.cpp`, which already handles `IDMap`, and round-trips with `index_factory` for IDMap-wrapped indexes such as `IDMap,Flat` and `IDMap,IVF100,Flat`.

Reviewed By: mnorris11

Differential Revision: D107159123

fbshipit-source-id: 40dca17a78743bf53373ec3077fd640e0fff76d4

contrib/factory_tools.py

@@ -146,4 +146,11 @@ def reverse_index_factory(index):
         }
         return f"SQ{sqtypes[index.sq.qtype]}"
 
+    # IndexIDMap2 is a subclass of IndexIDMap, so it must be checked first.
+    elif isinstance(index, faiss.IndexIDMap2):
+        return f"IDMap2,{reverse_index_factory(index.index)}"
+
+    elif isinstance(index, faiss.IndexIDMap):
+        return f"IDMap,{reverse_index_factory(index.index)}"
+
     raise NotImplementedError()

tests/test_contrib.py

@@ -21,6 +21,7 @@ from faiss.contrib import (
     clustering,
     datasets,
     evaluation,
+    factory_tools,
     inspect_tools,
     ivf_tools,
 )
@@ -834,3 +835,17 @@ class TestMerge(unittest.TestCase):
     def test_ondisk_merge_with_shift_ids(self):
         # verified that recall is same for test_ondisk_merge and
         self.do_test_ondisk_merge(True)
+
+
+class TestFactoryTools(unittest.TestCase):
+
+    def test_idmap(self):
+        index = faiss.IndexIDMap(faiss.IndexFlatL2(32))
+        self.assertEqual(
+            factory_tools.reverse_index_factory(index), "IDMap,Flat"
+        )
+        # IndexIDMap2 subclasses IndexIDMap, so it must be matched first.
+        index2 = faiss.IndexIDMap2(faiss.IndexFlatL2(32))
+        self.assertEqual(
+            factory_tools.reverse_index_factory(index2), "IDMap2,Flat"
+        )