diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index dcee6fee2..01e1b95af 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -1110,6 +1110,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_ProductQuantizer(&idxp->pq, f); idxp->code_size = idxp->pq.code_size; read_vector(idxp->codes, f); + FAISS_THROW_IF_NOT( + idxp->codes.size() == idxp->ntotal * idxp->code_size); if (h == fourcc("IxPo") || h == fourcc("IxPq")) { READ1(idxp->search_type); READ1(idxp->encode_signs); @@ -1132,6 +1134,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { } READ1(idxr->code_size); read_vector(idxr->codes, f); + FAISS_THROW_IF_NOT( + idxr->codes.size() == idxr->ntotal * idxr->code_size); idx = std::move(idxr); } else if (h == fourcc("IxLS")) { auto idxr = std::make_unique(); @@ -1139,6 +1143,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_LocalSearchQuantizer(idxr->lsq, f); READ1(idxr->code_size); read_vector(idxr->codes, f); + FAISS_THROW_IF_NOT( + idxr->codes.size() == idxr->ntotal * idxr->code_size); idx = std::move(idxr); } else if (h == fourcc("IxPR")) { auto idxpr = std::make_unique(); @@ -1146,6 +1152,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_ProductResidualQuantizer(idxpr->prq, f, io_flags); READ1(idxpr->code_size); read_vector(idxpr->codes, f); + FAISS_THROW_IF_NOT( + idxpr->codes.size() == idxpr->ntotal * idxpr->code_size); idx = std::move(idxpr); } else if (h == fourcc("IxPL")) { auto idxpl = std::make_unique(); @@ -1153,6 +1161,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_ProductLocalSearchQuantizer(idxpl->plsq, f); READ1(idxpl->code_size); read_vector(idxpl->codes, f); + FAISS_THROW_IF_NOT( + idxpl->codes.size() == idxpl->ntotal * idxpl->code_size); idx = std::move(idxpl); } else if (h == fourcc("ImRQ")) { auto idxr = std::make_unique(); diff --git a/tests/test_read_index_deserialize.cpp b/tests/test_read_index_deserialize.cpp index 4cd9be08d..22f3dd099 100644 --- a/tests/test_read_index_deserialize.cpp +++ b/tests/test_read_index_deserialize.cpp @@ -938,3 +938,165 @@ TEST(ReadIndexDeserialize, SQTrainedSizeMismatch) { expect_read_throws_with(buf, "ScalarQuantizer trained size"); } + +// ----------------------------------------------------------------------- +// Test: IndexPQ codes vector size mismatch. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexPQCodesSizeMismatch) { + // "IxPq": fourcc + index_header + PQ + codes + search_type + + // encode_signs + polysemous_ht + std::vector buf; + push_fourcc(buf, "IxPq"); + push_index_header(buf, /*d=*/4, /*ntotal=*/2); + // PQ: d=4, M=2, nbits=8 -> code_size=2, ksub=256 + std::vector centroids(4 * 256, 0.0f); + push_pq(buf, /*d=*/4, /*M=*/2, /*nbits=*/8, centroids); + // codes: should be ntotal * code_size = 2 * 2 = 4 bytes + // Provide wrong size (10 bytes) + push_vector(buf, std::vector(10, 0)); + // IxPq format extras: + push_val(buf, 0); // search_type + push_val(buf, 0); // encode_signs + push_val(buf, 0); // polysemous_ht + + expect_read_throws_with(buf, "codes.size()"); +} + +/// Helper: append a minimal AdditiveQuantizer (d, M, nbits, is_trained, +/// codebooks, search_type=ST_decompress, norm_min, norm_max). +static void push_additive_quantizer( + std::vector& buf, + size_t d, + size_t M, + const std::vector& nbits) { + push_val(buf, d); + push_val(buf, M); + push_vector(buf, nbits); + push_val(buf, true); // is_trained + push_vector(buf, {}); // codebooks (empty) + push_val(buf, 0); // search_type = ST_decompress + push_val(buf, 0.0f); // norm_min + push_val(buf, 1.0f); // norm_max +} + +/// Helper: append a minimal ResidualQuantizer (AdditiveQuantizer + +/// train_type + max_beam_size). Uses Skip_codebook_tables to avoid +/// compute_codebook_tables on empty codebooks. +static void push_residual_quantizer( + std::vector& buf, + size_t d, + size_t M, + const std::vector& nbits) { + push_additive_quantizer(buf, d, M, nbits); + push_val(buf, 2048); // train_type = Skip_codebook_tables + push_val(buf, 1); // max_beam_size +} + +/// Helper: append a minimal LocalSearchQuantizer (AdditiveQuantizer + +/// K + train_iters + encode_ils_iters + train_ils_iters + icm_iters + +/// p + lambd + chunk_size + random_seed + nperts + +/// update_codebooks_with_double). +static void push_local_search_quantizer( + std::vector& buf, + size_t d, + size_t M, + const std::vector& nbits) { + push_additive_quantizer(buf, d, M, nbits); + push_val(buf, 256); // K + push_val(buf, 25); // train_iters + push_val(buf, 8); // encode_ils_iters + push_val(buf, 8); // train_ils_iters + push_val(buf, 4); // icm_iters + push_val(buf, 0.5f); // p + push_val(buf, 0.0f); // lambd + push_val(buf, 0); // chunk_size + push_val(buf, 123); // random_seed + push_val(buf, 4); // nperts + push_val(buf, false); // update_codebooks_with_double +} + +// ----------------------------------------------------------------------- +// Test: IndexResidualQuantizer codes vector size mismatch. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexResidualQuantizerCodesSizeMismatch) { + // "IxRq": fourcc + index_header + ResidualQuantizer + code_size + codes + std::vector buf; + push_fourcc(buf, "IxRq"); + push_index_header(buf, /*d=*/4, /*ntotal=*/2); + push_residual_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8}); + // code_size = 1 (M=1, nbits=8 → 1 byte per code) + push_val(buf, 1); // code_size + // codes: should be ntotal * code_size = 2 * 1 = 2 bytes + // Provide wrong size (10 bytes) + push_vector(buf, std::vector(10, 0)); + + expect_read_throws_with(buf, "codes.size()"); +} + +// ----------------------------------------------------------------------- +// Test: IndexLocalSearchQuantizer codes vector size mismatch. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexLocalSearchQuantizerCodesSizeMismatch) { + // "IxLS": fourcc + index_header + LocalSearchQuantizer + code_size + codes + std::vector buf; + push_fourcc(buf, "IxLS"); + push_index_header(buf, /*d=*/4, /*ntotal=*/2); + push_local_search_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8}); + // code_size = 1 (M=1, nbits=8 → 1 byte per code) + push_val(buf, 1); // code_size + // codes: should be ntotal * code_size = 2 * 1 = 2 bytes + // Provide wrong size (10 bytes) + push_vector(buf, std::vector(10, 0)); + + expect_read_throws_with(buf, "codes.size()"); +} + +// ----------------------------------------------------------------------- +// Test: IndexProductResidualQuantizer codes vector size mismatch. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexProductResidualQuantizerCodesSizeMismatch) { + // "IxPR": fourcc + index_header + ProductResidualQuantizer + code_size + // + codes + // ProductResidualQuantizer = ProductAdditiveQuantizer(AdditiveQuantizer + // + nsplits) + nsplits * ResidualQuantizer + std::vector buf; + push_fourcc(buf, "IxPR"); + push_index_header(buf, /*d=*/4, /*ntotal=*/2); + // ProductAdditiveQuantizer: AdditiveQuantizer + nsplits + push_additive_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8}); + push_val(buf, 1); // nsplits = 1 + // 1 nested ResidualQuantizer (sub-dimension = d / nsplits = 4) + push_residual_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8}); + // code_size = 1 + push_val(buf, 1); // code_size + // codes: should be ntotal * code_size = 2 * 1 = 2 bytes + // Provide wrong size (10 bytes) + push_vector(buf, std::vector(10, 0)); + + expect_read_throws_with(buf, "codes.size()"); +} + +// ----------------------------------------------------------------------- +// Test: IndexProductLocalSearchQuantizer codes vector size mismatch. +// ----------------------------------------------------------------------- +TEST(ReadIndexDeserialize, IndexProductLocalSearchQuantizerCodesSizeMismatch) { + // "IxPL": fourcc + index_header + ProductLocalSearchQuantizer + code_size + // + codes + // ProductLocalSearchQuantizer = ProductAdditiveQuantizer(AdditiveQuantizer + // + nsplits) + nsplits * LocalSearchQuantizer + std::vector buf; + push_fourcc(buf, "IxPL"); + push_index_header(buf, /*d=*/4, /*ntotal=*/2); + // ProductAdditiveQuantizer: AdditiveQuantizer + nsplits + push_additive_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8}); + push_val(buf, 1); // nsplits = 1 + // 1 nested LocalSearchQuantizer (sub-dimension = d / nsplits = 4) + push_local_search_quantizer(buf, /*d=*/4, /*M=*/1, /*nbits=*/{8}); + // code_size = 1 + push_val(buf, 1); // code_size + // codes: should be ntotal * code_size = 2 * 1 = 2 bytes + // Provide wrong size (10 bytes) + push_vector(buf, std::vector(10, 0)); + + expect_read_throws_with(buf, "codes.size()"); +}