Fix reverse_index_factory formatting of ScalarQuantizers (#4003)

Summary: Pull Request resolved: #4003 It was emitting "SQSQ8" for `IndexScalarQuantizer`. The enum names already have `SQ` in the prefix. Reviewed By: kuarora Differential Revision: D65083146 fbshipit-source-id: 564f6d3d8987a2188fb7aa82b5ba772034225550
facebookresearch · Oct 29, 2024 · c5fbefe · c5fbefe
1 parent acaa01f
commit c5fbefe
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 46 deletions.
diff --git a/faiss/cppcontrib/factory_tools.cpp b/faiss/cppcontrib/factory_tools.cpp
@@ -143,7 +143,7 @@ std::string reverse_index_factory(const faiss::Index* index) {
     } else if (
             const faiss::IndexScalarQuantizer* sq_index =
                     dynamic_cast<const faiss::IndexScalarQuantizer*>(index)) {
-        return std::string("SQ") + sq_types.at(sq_index->sq.qtype);
+        return sq_types.at(sq_index->sq.qtype);
     } else if (
             const faiss::IndexIDMap* idmap =
                     dynamic_cast<const faiss::IndexIDMap*>(index)) {

diff --git a/tests/test_factory_tools.cpp b/tests/test_factory_tools.cpp
@@ -9,51 +9,34 @@
 #include <faiss/index_factory.h>
 #include <gtest/gtest.h>
 
-using namespace faiss;
+namespace faiss {
 
 TEST(TestFactoryTools, TestReverseIndexFactory) {
-    auto factory_string = "Flat";
-    auto index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "IMI2x5,PQ8x8";
-    index = faiss::index_factory(32, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "IVF32_HNSW32,SQ8";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "IVF8,Flat";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "IVF8,SQ4";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "IVF8,PQ4x8";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "LSHrt";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "PQ4x8";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
-
-    factory_string = "HNSW32";
-    index = faiss::index_factory(64, factory_string);
-    EXPECT_EQ(factory_string, reverse_index_factory(index));
-    delete index;
+    for (const char* factory : {
+                 "Flat",
+                 "IMI2x5,PQ8x8",
+                 "IVF32_HNSW32,SQ8",
+                 "IVF8,Flat",
+                 "IVF8,SQ4",
+                 "IVF8,PQ4x8",
+                 "LSHrt",
+                 "PQ4x8",
+                 "HNSW32",
+                 "SQ8",
+                 "SQfp16",
+         }) {
+        std::unique_ptr<Index> index{index_factory(64, factory)};
+        ASSERT_TRUE(index);
+        EXPECT_EQ(factory, reverse_index_factory(index.get()));
+    }
+    using Case = std::pair<const char*, const char*>;
+    for (auto [src, dst] : {
+                 Case{"SQ8,RFlat", "SQ8,Refine(Flat)"},
+         }) {
+        std::unique_ptr<Index> index{index_factory(64, src)};
+        ASSERT_TRUE(index);
+        EXPECT_EQ(dst, reverse_index_factory(index.get()));
+    }
 }
+
+} // namespace faiss