facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Validate AdditiveQuantizer dimension consistency during deserialization (#5023)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5023

Add deserialization-time validation to catch AdditiveQuantizer dimension
mismatches that would cause crashes during search operations.

1. **AQ dimension vs index dimension cross-check**:
   Added `validate_aq_dimension_match()` helper that verifies `aq.d == idx.d`
   for all AdditiveQuantizer-based index types. Without this check, a
   corrupted index where AQ.d != Index.d causes `compute_LUT()` to read
   codebooks out of bounds (query vectors are sized for Index.d while
   codebooks are laid out for AQ.d).

2. **Coverage across all AQ index types**:
   - IndexResidualQuantizer (IxRq/IxRQ)
   - IndexLocalSearchQuantizer (IxLS)
   - IndexProductResidualQuantizer (IxPR)
   - IndexProductLocalSearchQuantizer (IxPL)
   - ResidualCoarseQuantizer (ImRQ)
   - IndexAdditiveQuantizerFastScan (IRfs/ILfs/IPRf/IPLf)
   - IndexIVFAdditiveQuantizerFastScan (IVLf/IVRf/NPLf/NPRf)
   - IndexIVFAdditiveQuantizer (IwLS/IwRQ/IwPL/IwPR)

Reviewed By: mnorris11

Differential Revision: D99137777

fbshipit-source-id: bb4201f32b1268d66907f9a3f34eff869c80dc9c
This commit is contained in:
Justin Gibbs
2026-04-03 09:58:37 -07:00
committed by meta-codesync[bot]
parent 553c78b256
commit 408786bec9
2 changed files with 84 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/impl/index_read.cpp
+32 -25
View File
@@ -697,6 +697,22 @@ static void validate_codebooks_size(
aq.total_codebook_size);
}
// Validate that the AdditiveQuantizer dimension matches the index header
// dimension. compute_LUT() treats codebooks as a (d, total_codebook_size)
// matrix and query vectors are sized for idx_d, so a mismatch leads to
// out-of-bounds reads.
static void validate_aq_dimension_match(
const AdditiveQuantizer& aq,
int idx_d,
const char* index_type) {
FAISS_THROW_IF_NOT_FMT(
aq.d == static_cast<size_t>(idx_d),
"%s: AdditiveQuantizer d=%zd does not match index d=%d",
index_type,
aq.d,
idx_d);
}
static void read_ResidualQuantizer(
ResidualQuantizer& rq,
IOReader* f,
@@ -1359,11 +1375,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
} else {
read_ResidualQuantizer(idxr->rq, f, io_flags);
}
FAISS_THROW_IF_NOT_FMT(
idxr->rq.d == static_cast<size_t>(idxr->d),
"IndexResidualQuantizer d mismatch: rq.d=%zd vs idx.d=%d",
idxr->rq.d,
idxr->d);
validate_aq_dimension_match(
idxr->rq, idxr->d, "IndexResidualQuantizer");
READ1(idxr->code_size);
read_vector(idxr->codes, f);
FAISS_THROW_IF_NOT(
@@ -1373,11 +1386,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
auto idxr = std::make_unique<IndexLocalSearchQuantizer>();
read_index_header(*idxr, f);
read_LocalSearchQuantizer(idxr->lsq, f);
FAISS_THROW_IF_NOT_FMT(
idxr->lsq.d == static_cast<size_t>(idxr->d),
"IndexLocalSearchQuantizer d mismatch: lsq.d=%zd vs idx.d=%d",
idxr->lsq.d,
idxr->d);
validate_aq_dimension_match(
idxr->lsq, idxr->d, "IndexLocalSearchQuantizer");
READ1(idxr->code_size);
read_vector(idxr->codes, f);
FAISS_THROW_IF_NOT(
@@ -1387,11 +1397,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
auto idxpr = std::make_unique<IndexProductResidualQuantizer>();
read_index_header(*idxpr, f);
read_ProductResidualQuantizer(idxpr->prq, f, io_flags);
FAISS_THROW_IF_NOT_FMT(
idxpr->prq.d == static_cast<size_t>(idxpr->d),
"IndexProductResidualQuantizer d mismatch: prq.d=%zd vs idx.d=%d",
idxpr->prq.d,
idxpr->d);
validate_aq_dimension_match(
idxpr->prq, idxpr->d, "IndexProductResidualQuantizer");
READ1(idxpr->code_size);
read_vector(idxpr->codes, f);
FAISS_THROW_IF_NOT(
@@ -1401,11 +1408,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
auto idxpl = std::make_unique<IndexProductLocalSearchQuantizer>();
read_index_header(*idxpl, f);
read_ProductLocalSearchQuantizer(idxpl->plsq, f);
FAISS_THROW_IF_NOT_FMT(
idxpl->plsq.d == static_cast<size_t>(idxpl->d),
"IndexProductLocalSearchQuantizer d mismatch: plsq.d=%zd vs idx.d=%d",
idxpl->plsq.d,
idxpl->d);
validate_aq_dimension_match(
idxpl->plsq, idxpl->d, "IndexProductLocalSearchQuantizer");
READ1(idxpl->code_size);
read_vector(idxpl->codes, f);
FAISS_THROW_IF_NOT(
@@ -1415,11 +1419,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
auto idxr = std::make_unique<ResidualCoarseQuantizer>();
read_index_header(*idxr, f);
read_ResidualQuantizer(idxr->rq, f, io_flags);
FAISS_THROW_IF_NOT_FMT(
idxr->rq.d == static_cast<size_t>(idxr->d),
"ResidualCoarseQuantizer d mismatch: rq.d=%zd vs idx.d=%d",
idxr->rq.d,
idxr->d);
validate_aq_dimension_match(
idxr->rq, idxr->d, "ResidualCoarseQuantizer");
READ1(idxr->beam_factor);
if (io_flags & IO_FLAG_SKIP_PRECOMPUTE_TABLE) {
// then we force the beam factor to -1
@@ -1490,6 +1491,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_ProductResidualQuantizer(
*(ProductResidualQuantizer*)idxaqfs->aq, f, io_flags);
}
validate_aq_dimension_match(
*idxaqfs->aq, idxaqfs->d, "IndexAdditiveQuantizerFastScan");
READ1(idxaqfs->implem);
READ1(idxaqfs->bbs);
@@ -1563,6 +1566,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_ProductResidualQuantizer(
*(ProductResidualQuantizer*)ivaqfs->aq, f, io_flags);
}
validate_aq_dimension_match(
*ivaqfs->aq, ivaqfs->d, "IndexIVFAdditiveQuantizerFastScan");
READ1(ivaqfs->by_residual);
READ1(ivaqfs->implem);
@@ -1720,6 +1725,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
read_ProductResidualQuantizer(
*(ProductResidualQuantizer*)iva->aq, f, io_flags);
}
validate_aq_dimension_match(
*iva->aq, iva->d, "IndexIVFAdditiveQuantizer");
READ1(iva->by_residual);
READ1(iva->use_precomputed_table);
read_InvertedLists(*iva, f, io_flags);
tests/test_read_index_deserialize.cpp
+52
View File
@@ -2868,3 +2868,55 @@ TEST(ReadIndexDeserialize, SQUntrainedEmptyTrainedAccepted) {
auto idx = read_index_up(&reader);
EXPECT_FALSE(idx->is_trained);
}
// Test: IndexResidualQuantizer rejects AQ dimension != index dimension.
TEST(ReadIndexDeserialize, IndexResidualQuantizerAQDimensionMismatch) {
std::vector<uint8_t> buf;
push_fourcc(buf, "IxRq");
// Index header says d=4, but AQ will have d=8 → mismatch.
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
push_residual_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4});
push_val<size_t>(buf, 1); // code_size
push_vector<uint8_t>(buf, {}); // codes
expect_read_throws_with(buf, "does not match index d");
}
// Test: IndexLocalSearchQuantizer rejects AQ dimension != index dimension.
TEST(ReadIndexDeserialize, IndexLocalSearchQuantizerAQDimensionMismatch) {
std::vector<uint8_t> buf;
push_fourcc(buf, "IxLS");
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
push_local_search_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4});
push_val<size_t>(buf, 1); // code_size
push_vector<uint8_t>(buf, {}); // codes
expect_read_throws_with(buf, "does not match index d");
}
// Test: ResidualCoarseQuantizer rejects AQ dimension != index dimension.
TEST(ReadIndexDeserialize, ResidualCoarseQuantizerAQDimensionMismatch) {
std::vector<uint8_t> buf;
push_fourcc(buf, "ImRQ");
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
push_residual_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4});
push_val<float>(buf, -1.0f); // beam_factor
expect_read_throws_with(buf, "does not match index d");
}
// Test: IndexResidualQuantizerFastScan rejects AQ dimension != index dimension.
TEST(ReadIndexDeserialize, IndexRQFastScanAQDimensionMismatch) {
std::vector<uint8_t> buf;
push_fourcc(buf, "IRfs");
// Index header says d=4, but AQ will have d=8 → mismatch.
push_index_header(buf, /*d=*/4, /*ntotal=*/0);
push_residual_quantizer(buf, /*d=*/8, /*M=*/1, /*nbits=*/{4});
// FastScan fields won't be reached due to early validation.
// But include a few just in case:
push_val<int>(buf, 0); // implem
push_val<int>(buf, 32); // bbs
push_val<int>(buf, 0); // qbs
expect_read_throws_with(buf, "does not match index d");
}