Files
faiss/tests/test_factory_tools.cpp
Alibek Zhakubayev e1f8d4fc0d Fix reverse_index_factory for multi-bit RaBitQ indexes (#5419)
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5419

### THIS DIFF

Fix two bugs in `reverse_index_factory` for multi-bit RaBitQ indexes, in
both the C++ and Python implementations.

**C++ (`factory_tools.cpp`):**
- `IndexIVFRaBitQ` case returned `prefix + ",RaBitQ"` unconditionally,
  regardless of `rabitq.nb_bits`. For a 2-bit index built from
  `"IVF_nlist,RaBitQ2"`, the function returned `"IVF_nlist,RaBitQ"` (1-bit
  factory string). This means `TelemetryWrapper` logs the wrong index type
  for every production multi-bit RaBitQ IVF index.
- `IndexRaBitQ` (standalone) had the same flaw — always returned
  `"RaBitQ"` regardless of `nb_bits`.

**Python (`contrib/factory_tools.py`):**
- Neither `IndexIVFRaBitQ` nor `IndexRaBitQ` had a case in
  `reverse_index_factory`, so any Python caller using this API on a RaBitQ
  index raised `NotImplementedError`.

**Fix:** Read `rabitq.nb_bits` from the typed cast result and append the
bit count when `nb_bits > 1` (e.g., `"RaBitQ2"`). For `nb_bits == 1`
the output is still `"RaBitQ"`, preserving backward compatibility with
existing logs and round-trip tests.

### CONTEXT

Multi-bit RaBitQ (`nb_bits=2..9`) was added alongside standard 1-bit
RaBitQ. `index_factory` has parsed `"RaBitQ2"`, `"RaBitQ4"`, etc., since
the feature landed, but `reverse_index_factory` was never updated to emit
these strings. Every production `IndexIVFRaBitQ` with `nb_bits > 1`
(including the RaBitQ2 deployment) logs `",RaBitQ"` (1-bit) in telemetry.
The Python API has silently been broken for all RaBitQ index types.

The fix pattern matches how other parameterized types (`PQMxN`, `HNSWNM`,
etc.) embed their parameters in the factory string.

Reviewed By: mnorris11

Differential Revision: D110042828

fbshipit-source-id: d4e976bd9fc99691a20fb33d893af2d36deeb1e1
2026-07-14 12:27:43 -07:00

79 lines
2.3 KiB
C++

/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <faiss/IndexIDMap.h>
#include <faiss/factory_tools.h>
#include <faiss/index_factory.h>
#include <gtest/gtest.h>
namespace faiss {
TEST(TestFactoryTools, TestReverseIndexFactory) {
for (const char* factory : {
"Flat",
"IMI2x5,PQ8x8",
"IVF32_HNSW32,SQ8",
"IVF8,Flat",
"IVF8,SQ4",
"IVF8,PQ4x8",
"LSHrt",
"PQ4x8",
"HNSW32",
"SQ8",
"SQtqmse1",
"SQtqmse2",
"SQtqmse3",
"SQtqmse4",
"SQtqmse8",
"SQfp16",
"NSG24,Flat",
"NSG16,SQ8",
"RaBitQ",
"RaBitQ2",
"RaBitQ4",
"IVF8,RaBitQ",
"IVF8,RaBitQ2",
"IVF8,RaBitQ4",
"IVF8,SQtqmse8",
"IVF8,SQ0",
"EDEN",
"EDEN4",
"EDEN4BIASED",
"IVF8,EDEN",
"IVF8,EDEN4",
"IVF8,EDEN4BIASED",
}) {
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)"},
Case{"NSG", "NSG32,Flat"},
Case{"NSG,PQ8", "NSG32,PQ8x8"},
}) {
std::unique_ptr<Index> index{index_factory(64, src)};
ASSERT_TRUE(index);
EXPECT_EQ(dst, reverse_index_factory(index.get()));
}
}
// IndexIDMap2 inherits IndexIDMap; verify the subclass is distinguished.
TEST(TestFactoryTools, TestReverseIndexFactoryIDMap) {
std::unique_ptr<Index> base{index_factory(64, "Flat")};
ASSERT_TRUE(base);
IndexIDMap idmap(base.get());
EXPECT_EQ("IDMap,Flat", reverse_index_factory(&idmap));
IndexIDMap2 idmap2(base.get());
EXPECT_EQ("IDMap2,Flat", reverse_index_factory(&idmap2));
}
} // namespace faiss