From b4c66ba6031d1f659163e43430f8ce745b0dbe79 Mon Sep 17 00:00:00 2001 From: kzahiri1 Date: Fri, 4 Sep 2026 11:14:06 -0700 Subject: [PATCH] Fix reverse_index_factory/get_code_size round-trip for HNSW, IMI and IVFPQR (#5564) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Three defects in `contrib/factory_tools.py` where a reversed factory string rebuilds a different index or fails to parse. `test_get_code_size_hnsw_roundtrip` already pins this invariant for flat HNSW; these cases escape it. 1. **HNSW storage dropped.** The `IndexHNSW` branch never read `index.storage`, so `HNSW32,SQ8` and `HNSW32,PQ16x8` both reversed to `"HNSW32"` and rebuilt as `IndexHNSWFlat`. `get_code_size` then applied the flat `d*4` formula: at d=128 it reported 768 bytes/vector for an index whose storage is 128 plus the 256-byte graph estimate, so 384. 2. **IMI regex fields transposed.** `index_factory.cpp:295` parses `IMI2x([0-9]+)`, but the regex read `IMI\d+x2`. Every real IMI key (`IMI2x10,PQ16`, and the ones in `demos/demo_auto_tune.py`) raised `RuntimeError("cannot parse")`. 3. **IndexIVFPQR matched by the IndexIVFPQ check**, dropping the refinement stage: `IVF64,PQ8+16` reversed to `IVF64,PQ8x8`, 8 bytes instead of 24. Consumers are `benchs/bench_all_ivf/parse_bench_all_ivf.py` (imports `get_code_size` as `unitsize`) and `benchs/bench_fw/index.py`. Flat storage stays implicit as the bare `HNSW` string, and a null `storage` pointer still returns `HNSW` instead of raising. **Behavior change worth flagging:** `IndexHNSW2Level` storage is `Index2Layer`, which `reverse_index_factory` does not handle, so it now raises `NotImplementedError` instead of silently returning a wrong `"HNSW32"`. It is not reachable from `index_factory`. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5564 Test Plan: Four tests added to `TestFactoryTools`. All four fail on unpatched `factory_tools.py` and pass with it: ``` # pristine contrib/factory_tools.py, new tests present 4 failed, 7 passed FAILED test_get_code_size_hnsw_non_flat_storage FAILED test_get_code_size_imi FAILED test_hnsw_storage_reverse_index_factory FAILED test_ivfpqr_reverse_index_factory # with the fix 11 passed ``` Full `tests/test_contrib.py tests/test_factory.py`, pristine vs branch: ``` pristine: 2 failed, 88 passed branch: 2 failed, 92 passed ``` Failure sets are identical. The two are `TestComputeGT::test_compute_GT_gpu` and `test_compute_GT_ip_gpu`, pre-existing and unrelated (`AttributeError: module 'faiss' has no attribute 'GpuMultipleClonerOptions'` on a CPU-only build). An A/B sweep over 40 factory strings changes exactly the 7 expected lines and nothing else: ``` - HNSW32,SQ8 IndexHNSWSQ HNSW32 768 + HNSW32,SQ8 IndexHNSWSQ HNSW32,SQ8 384 - HNSW32,SQfp16 IndexHNSWSQ HNSW32 768 + HNSW32,SQfp16 IndexHNSWSQ HNSW32,SQfp16 512 - HNSW32,PQ16x8 IndexHNSWPQ HNSW32 768 + HNSW32,PQ16x8 IndexHNSWPQ HNSW32,PQ16x8 272 - IVF64,PQ8+16 IndexIVFPQR IVF64,PQ8x8 8 + IVF64,PQ8+16 IndexIVFPQR IVF64,PQ8+16 24 - IMI2x5,Flat IndexIVFFlat IMI2x5,Flat + IMI2x5,Flat IndexIVFFlat IMI2x5,Flat 512 - IMI2x5,PQ8 IndexIVFPQ IMI2x5,PQ8x8 + IMI2x5,PQ8 IndexIVFPQ IMI2x5,PQ8x8 8 - IMI2x8,PQ8x8 IndexIVFPQ IMI2x8,PQ8x8 + IMI2x8,PQ8x8 IndexIVFPQ IMI2x8,PQ8x8 8 ``` flake8 at the 80-char limit from CONTRIBUTING reports the same 12 pre-existing findings before and after, none in the changed lines. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Reviewed By: mnorris11 Differential Revision: D118661065 Pulled By: trang-nm-nguyen fbshipit-source-id: b50a9522361e3fd14f276b36196d3c728567b980 --- contrib/factory_tools.py | 17 ++++++++++++-- tests/test_contrib.py | 50 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/contrib/factory_tools.py b/contrib/factory_tools.py index 1470d3433..f70a85e00 100644 --- a/contrib/factory_tools.py +++ b/contrib/factory_tools.py @@ -25,7 +25,7 @@ def get_code_size(d, indexkey): if mo: return get_code_size(d, mo.group(1)) - mo = re.match("IMI\\d+x2,(.*)$", indexkey) + mo = re.match("IMI\\d+x\\d+,(.*)$", indexkey) if mo: return get_code_size(d, mo.group(1)) @@ -50,6 +50,11 @@ def get_code_size(d, indexkey): M = int(mo.group(1)) return d * 4 + M * 2 * 4 # roughly + mo = re.match("HNSW(\\d+),(.*)$", indexkey) + if mo: + M = int(mo.group(1)) + return get_code_size(d, mo.group(2)) + M * 2 * 4 # roughly + if indexkey == "SQ8": return d elif indexkey == "SQ4": @@ -127,6 +132,9 @@ def reverse_index_factory(index): return prefix + ",Flat" if isinstance(index, faiss.IndexIVFScalarQuantizer): return prefix + "," + sq_names[index.sq.qtype] + # IndexIVFPQR subclasses IndexIVFPQ, so it must be checked first. + if isinstance(index, faiss.IndexIVFPQR): + return prefix + f",PQ{index.pq.M}+{index.refine_pq.M}" if isinstance(index, faiss.IndexIVFPQ): return prefix + f",PQ{index.pq.M}x{index.pq.nbits}" if isinstance(index, faiss.IndexIVFPQFastScan): @@ -152,7 +160,12 @@ def reverse_index_factory(index): return f"{prefix},{reverse_index_factory(index.index)}" elif isinstance(index, faiss.IndexHNSW): - return f"HNSW{get_hnsw_M(index)}" + prefix = f"HNSW{get_hnsw_M(index)}" + storage = faiss.downcast_index(index.storage) + # flat storage is implicit in the bare HNSW factory string + if storage is None or isinstance(storage, faiss.IndexFlat): + return prefix + return f"{prefix},{reverse_index_factory(storage)}" elif isinstance(index, faiss.IndexRefine): return ( diff --git a/tests/test_contrib.py b/tests/test_contrib.py index 6d363684c..7d5faeda1 100644 --- a/tests/test_contrib.py +++ b/tests/test_contrib.py @@ -1021,6 +1021,56 @@ class TestFactoryTools(unittest.TestCase): code_size = factory_tools.get_code_size(d, factory_str) self.assertEqual(code_size, d * 4 + 16 * 2 * 4) + def test_hnsw_storage_reverse_index_factory(self): + d = 128 + # flat storage stays implicit, everything else must be spelled out + cases = { + "HNSW32,Flat": "HNSW32", + "HNSW32,SQ8": "HNSW32,SQ8", + "HNSW32,SQfp16": "HNSW32,SQfp16", + "HNSW32,PQ16x8": "HNSW32,PQ16x8", + } + for key, expected in cases.items(): + index = faiss.index_factory(d, key) + factory_str = factory_tools.reverse_index_factory(index) + self.assertEqual(factory_str, expected) + # the string must rebuild the same index class + rebuilt = faiss.index_factory(d, factory_str) + self.assertEqual(type(rebuilt), type(index)) + + def test_get_code_size_hnsw_non_flat_storage(self): + d = 128 + graph = 32 * 2 * 4 + self.assertEqual( + factory_tools.get_code_size(d, "HNSW32,SQ8"), d + graph + ) + self.assertEqual( + factory_tools.get_code_size(d, "HNSW32,SQfp16"), d * 2 + graph + ) + self.assertEqual( + factory_tools.get_code_size(d, "HNSW32,PQ16x8"), 16 + graph + ) + + def test_get_code_size_imi(self): + # the coarse quantizer adds no per-vector bytes + d = 64 + self.assertEqual(factory_tools.get_code_size(d, "IMI2x10,PQ16"), 16) + self.assertEqual(factory_tools.get_code_size(d, "IMI2x5,Flat"), d * 4) + # round-trip: get_code_size must parse what reverse emits + index = faiss.index_factory(d, "IMI2x5,PQ8") + factory_str = factory_tools.reverse_index_factory(index) + self.assertEqual(factory_str, "IMI2x5,PQ8x8") + self.assertEqual(factory_tools.get_code_size(d, factory_str), 8) + + def test_ivfpqr_reverse_index_factory(self): + d = 128 + index = faiss.index_factory(d, "IVF64,PQ8+16") + factory_str = factory_tools.reverse_index_factory(index) + self.assertEqual(factory_str, "IVF64,PQ8+16") + rebuilt = faiss.index_factory(d, factory_str) + self.assertEqual(type(rebuilt), type(index)) + self.assertEqual(factory_tools.get_code_size(d, factory_str), 8 + 16) + def test_rabitq_reverse_index_factory(self): d = 64 quantizer = faiss.IndexFlatL2(d)