faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/index_read_utils.h

@@ -19,7 +19,6 @@ namespace faiss {
 struct ProductQuantizer;
 struct ScalarQuantizer;
 
-void read_index_header(Index* idx, IOReader* f);
 void read_direct_map(DirectMap* dm, IOReader* f);
 void read_ivf_header(
         IndexIVF* ivf,
@@ -30,7 +29,7 @@ ArrayInvertedLists* set_array_invlist(
         IndexIVF* ivf,
         std::vector<std::vector<idx_t>>& ids);
 void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f);
-void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f);
+void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f, const Index& idx);
 
 } // namespace faiss
 

data/vendor/faiss/faiss/impl/index_write.cpp

@@ -12,10 +12,14 @@
 
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 
 #include <faiss/invlists/InvertedListsIOHook.h>
 
+#include <faiss/invlists/BlockInvertedLists.h>
+
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/RaBitQUtils.h>
 #include <faiss/utils/hamming.h>
 
 #include <faiss/Index2Layer.h>
@@ -49,6 +53,9 @@
 #ifdef FAISS_ENABLE_SVS
 #include <faiss/impl/svs_io.h>
 #include <faiss/svs/IndexSVSFlat.h>
+#include <faiss/svs/IndexSVSIVF.h>
+#include <faiss/svs/IndexSVSIVFLVQ.h>
+#include <faiss/svs/IndexSVSIVFLeanVec.h>
 #include <faiss/svs/IndexSVSVamana.h>
 #include <faiss/svs/IndexSVSVamanaLVQ.h>
 #include <faiss/svs/IndexSVSVamanaLeanVec.h>
@@ -101,7 +108,14 @@ static void write_index_header(const Index* idx, IOWriter* f) {
 }
 
 void write_VectorTransform(const VectorTransform* vt, IOWriter* f) {
-    if (const LinearTransform* lt = dynamic_cast<const LinearTransform*>(vt)) {
+    if (const HadamardRotation* hr =
+                dynamic_cast<const HadamardRotation*>(vt)) {
+        uint32_t h = fourcc("HRot");
+        WRITE1(h);
+        WRITE1(hr->seed);
+    } else if (
+            const LinearTransform* lt =
+                    dynamic_cast<const LinearTransform*>(vt)) {
         if (dynamic_cast<const RandomRotationMatrix*>(lt)) {
             uint32_t h = fourcc("rrot");
             WRITE1(h);
@@ -258,11 +272,20 @@ void write_InvertedLists(const InvertedLists* ils, IOWriter* f) {
     } else if (
             const auto& ailp =
                     dynamic_cast<const ArrayInvertedListsPanorama*>(ils)) {
-        uint32_t h = fourcc("ilpn");
-        WRITE1(h);
-        WRITE1(ailp->nlist);
-        WRITE1(ailp->code_size);
-        WRITE1(ailp->n_levels);
+        if (ailp->pano.batch_size == Panorama::kDefaultBatchSize) {
+            uint32_t h = fourcc("ilpn");
+            WRITE1(h);
+            WRITE1(ailp->nlist);
+            WRITE1(ailp->code_size);
+            WRITE1(ailp->n_levels);
+        } else {
+            uint32_t h = fourcc("ilp2");
+            WRITE1(h);
+            WRITE1(ailp->nlist);
+            WRITE1(ailp->code_size);
+            WRITE1(ailp->n_levels);
+            WRITE1(ailp->pano.batch_size);
+        }
         uint32_t list_type = fourcc("full");
         WRITE1(list_type);
         std::vector<size_t> sizes;
@@ -446,9 +469,9 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         uint32_t h = fourcc("null");
         WRITE1(h);
     } else if (
-            const IndexFlatL2Panorama* idxpan =
-                    dynamic_cast<const IndexFlatL2Panorama*>(idx)) {
-        uint32_t h = fourcc("IxFP");
+            const IndexFlatPanorama* idxpan =
+                    dynamic_cast<const IndexFlatPanorama*>(idx)) {
+        uint32_t h = fourcc(idxpan->metric_type == METRIC_L2 ? "IxFP" : "IxFp");
         WRITE1(h);
         WRITE1(idxpan->d);
         WRITE1(idxpan->n_levels);
@@ -637,13 +660,13 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
 
         write_InvertedLists(ivaqfs->invlists, f);
     } else if (
-            const ResidualCoarseQuantizer* idxr_2 =
+            const ResidualCoarseQuantizer* idxrcq =
                     dynamic_cast<const ResidualCoarseQuantizer*>(idx)) {
         uint32_t h = fourcc("ImRQ");
         WRITE1(h);
         write_index_header(idx, f);
-        write_ResidualQuantizer(&idxr_2->rq, f);
-        WRITE1(idxr_2->beam_factor);
+        write_ResidualQuantizer(&idxrcq->rq, f);
+        WRITE1(idxrcq->beam_factor);
     } else if (
             const Index2Layer* idxp_2 = dynamic_cast<const Index2Layer*>(idx)) {
         uint32_t h = fourcc("Ix2L");
@@ -696,10 +719,18 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
     } else if (
             const IndexIVFFlatPanorama* ivfp =
                     dynamic_cast<const IndexIVFFlatPanorama*>(idx)) {
-        uint32_t h = fourcc("IwPn");
-        WRITE1(h);
-        write_ivf_header(ivfp, f);
-        WRITE1(ivfp->n_levels);
+        if (ivfp->batch_size == Panorama::kDefaultBatchSize) {
+            uint32_t h = fourcc("IwPn");
+            WRITE1(h);
+            write_ivf_header(ivfp, f);
+            WRITE1(ivfp->n_levels);
+        } else {
+            uint32_t h = fourcc("IwP2");
+            WRITE1(h);
+            write_ivf_header(ivfp, f);
+            WRITE1(ivfp->n_levels);
+            WRITE1(ivfp->batch_size);
+        }
         write_InvertedLists(ivfp->invlists, f);
     } else if (
             const IndexIVFFlat* ivfl_2 =
@@ -937,13 +968,12 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
     } else if (
             const IndexRaBitQFastScan* idxqfs =
                     dynamic_cast<const IndexRaBitQFastScan*>(idx)) {
-        uint32_t h = fourcc("Irfs");
+        uint32_t h = fourcc("Irfn");
         WRITE1(h);
         write_index_header(idx, f);
         write_RaBitQuantizer(&idxqfs->rabitq, f);
         WRITEVECTOR(idxqfs->center);
         WRITE1(idxqfs->qb);
-        WRITEVECTOR(idxqfs->flat_storage);
         WRITE1(idxqfs->bbs);
         WRITE1(idxqfs->ntotal2);
         WRITE1(idxqfs->M2);
@@ -997,6 +1027,8 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
             h = fourcc("ILVQ"); // LVQ
         } else if (lean != nullptr) {
             h = fourcc("ISVL"); // LeanVec
+        } else if (svs->stored_vectors_valid && !svs->stored_vectors.empty()) {
+            h = fourcc("ISV2"); // uncompressed + stored_vectors
         } else {
             h = fourcc("ISVD"); // uncompressed
         }
@@ -1039,6 +1071,10 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
                 os.flush();
             }
         }
+
+        if (h == fourcc("ISV2")) {
+            WRITEVECTOR(svs->stored_vectors);
+        }
     } else if (
             const IndexSVSFlat* svs = dynamic_cast<const IndexSVSFlat*>(idx)) {
         uint32_t h = fourcc("ISVF");
@@ -1055,12 +1091,68 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
             svs->serialize_impl(os);
             os.flush();
         }
+    } else if (
+            const IndexSVSIVF* svs_ivf =
+                    dynamic_cast<const IndexSVSIVF*>(idx)) {
+        uint32_t h;
+        auto* lvq = dynamic_cast<const IndexSVSIVFLVQ*>(svs_ivf);
+        auto* lean = dynamic_cast<const IndexSVSIVFLeanVec*>(svs_ivf);
+        if (lvq != nullptr) {
+            h = fourcc("ISIQ"); // IndexSVSIVFLVQ
+        } else if (lean != nullptr) {
+            h = fourcc("ISIL"); // IndexSVSIVFLeanVec
+        } else {
+            h = fourcc("ISID"); // IndexSVSIVF
+        }
+
+        WRITE1(h);
+        write_index_header(svs_ivf, f);
+        WRITE1(svs_ivf->num_centroids);
+        WRITE1(svs_ivf->minibatch_size);
+        WRITE1(svs_ivf->num_iterations);
+        WRITE1(svs_ivf->is_hierarchical);
+        WRITE1(svs_ivf->training_fraction);
+        WRITE1(svs_ivf->hierarchical_level1_clusters);
+        WRITE1(svs_ivf->seed);
+        WRITE1(svs_ivf->n_probes);
+        WRITE1(svs_ivf->k_reorder);
+        WRITE1(svs_ivf->num_threads);
+        WRITE1(svs_ivf->intra_query_threads);
+        WRITE1(svs_ivf->storage_kind);
+        WRITE1(svs_ivf->is_static);
+
+        if (lean != nullptr) {
+            WRITE1(lean->leanvec_d);
+        }
+
+        bool initialized = (svs_ivf->impl != nullptr);
+        WRITE1(initialized);
+        if (initialized) {
+            faiss::BufferedIOWriter bwr(f);
+            faiss::svs_io::WriterStreambuf wbuf(&bwr);
+            std::ostream os(&wbuf);
+            svs_ivf->serialize_impl(os);
+            os.flush();
+        }
+
+        if (lean != nullptr) {
+            // Store training data info
+            bool trained = (lean->training_data != nullptr);
+            WRITE1(trained);
+            if (trained) {
+                faiss::BufferedIOWriter bwr(f);
+                faiss::svs_io::WriterStreambuf wbuf(&bwr);
+                std::ostream os(&wbuf);
+                lean->serialize_training_data(os);
+                os.flush();
+            }
+        }
     }
 #endif // FAISS_ENABLE_SVS
     else if (
             const IndexIVFRaBitQFastScan* ivrqfs =
                     dynamic_cast<const IndexIVFRaBitQFastScan*>(idx)) {
-        uint32_t h = fourcc("Iwrf");
+        uint32_t h = fourcc("Iwrn");
         WRITE1(h);
         write_ivf_header(ivrqfs, f);
         write_RaBitQuantizer(&ivrqfs->rabitq, f);
@@ -1072,7 +1164,6 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         WRITE1(ivrqfs->implem);
         WRITE1(ivrqfs->qb);
         WRITE1(ivrqfs->centered);
-        WRITEVECTOR(ivrqfs->flat_storage);
         write_InvertedLists(ivrqfs->invlists, f);
     } else {
         FAISS_THROW_MSG("don't know how to serialize this type of index");
@@ -1156,7 +1247,7 @@ static void write_binary_multi_hash_map(
         size_t ntotal,
         IOWriter* f) {
     int id_bits = 0;
-    while ((ntotal > ((idx_t)1 << id_bits))) {
+    while ((ntotal > (size_t(1) << id_bits))) {
         id_bits++;
     }
     WRITE1(id_bits);

data/vendor/faiss/faiss/impl/io.cpp

@@ -134,14 +134,14 @@ int FileIOWriter::filedescriptor() {
  * IO buffer
  ***********************************************************************/
 
-BufferedIOReader::BufferedIOReader(IOReader* reader, size_t bsz)
-        : reader(reader),
-          bsz(bsz),
+BufferedIOReader::BufferedIOReader(IOReader* reader_in, size_t bsz_in)
+        : reader(reader_in),
+          bsz(bsz_in),
           ofs(0),
           ofs2(0),
           b0(0),
           b1(0),
-          buffer(bsz) {}
+          buffer(bsz_in) {}
 
 size_t BufferedIOReader::operator()(void* ptr, size_t unitsize, size_t nitems) {
     size_t size = unitsize * nitems;
@@ -184,8 +184,8 @@ size_t BufferedIOReader::operator()(void* ptr, size_t unitsize, size_t nitems) {
     return nb / unitsize;
 }
 
-BufferedIOWriter::BufferedIOWriter(IOWriter* writer, size_t bsz)
-        : writer(writer), bsz(bsz), ofs2(0), b0(0), buffer(bsz) {}
+BufferedIOWriter::BufferedIOWriter(IOWriter* writer_in, size_t bsz_in)
+        : writer(writer_in), bsz(bsz_in), ofs2(0), b0(0), buffer(bsz_in) {}
 
 size_t BufferedIOWriter::operator()(
         const void* ptr,

data/vendor/faiss/faiss/impl/io_macros.h

@@ -17,6 +17,10 @@
  * always called f and thus is not passed in as a macro parameter.
  **************************************************************/
 
+namespace faiss {
+size_t get_deserialization_vector_byte_limit();
+} // namespace faiss
+
 #define READANDCHECK(ptr, n)                         \
     {                                                \
         size_t ret = (*f)(ptr, sizeof(*(ptr)), n);   \
@@ -37,14 +41,24 @@
         READ1(x);           \
     }
 
-// will fail if we write 256G of data at once...
-#define READVECTOR(vec)                                              \
-    {                                                                \
-        size_t size;                                                 \
-        READANDCHECK(&size, 1);                                      \
-        FAISS_THROW_IF_NOT(size >= 0 && size < (uint64_t{1} << 40)); \
-        (vec).resize(size);                                          \
-        READANDCHECK((vec).data(), size);                            \
+// Rejects vectors whose total allocation would exceed the configurable
+// byte limit (default 1 TB).
+#define READVECTOR(vec)                                                  \
+    {                                                                    \
+        size_t size;                                                     \
+        READANDCHECK(&size, 1);                                          \
+        FAISS_THROW_IF_NOT(                                              \
+                size >= 0 &&                                             \
+                size < (faiss::get_deserialization_vector_byte_limit() / \
+                        sizeof(*(vec).data())));                         \
+        FAISS_THROW_IF_NOT_FMT(                                          \
+                size <= SIZE_MAX / sizeof((vec)[0]),                     \
+                "READVECTOR: size %zu would overflow for element "       \
+                "size %zu",                                              \
+                size,                                                    \
+                sizeof((vec)[0]));                                       \
+        (vec).resize(size);                                              \
+        READANDCHECK((vec).data(), size);                                \
     }
 
 #define WRITEANDCHECK(ptr, n)                         \
@@ -78,12 +92,15 @@
         WRITEANDCHECK((vec).data(), size * 4);     \
     }
 
-#define READXBVECTOR(vec)                                            \
-    {                                                                \
-        size_t size;                                                 \
-        READANDCHECK(&size, 1);                                      \
-        FAISS_THROW_IF_NOT(size >= 0 && size < (uint64_t{1} << 40)); \
-        size *= 4;                                                   \
-        (vec).resize(size);                                          \
-        READANDCHECK((vec).data(), size);                            \
+#define READXBVECTOR(vec)                                                \
+    {                                                                    \
+        size_t size;                                                     \
+        READANDCHECK(&size, 1);                                          \
+        FAISS_THROW_IF_NOT(                                              \
+                size >= 0 &&                                             \
+                size < (faiss::get_deserialization_vector_byte_limit() / \
+                        (4 * sizeof(*(vec).data()))));                   \
+        size *= 4;                                                       \
+        (vec).resize(size);                                              \
+        READANDCHECK((vec).data(), size);                                \
     }

data/vendor/faiss/faiss/impl/kmeans1d.cpp

@@ -48,12 +48,12 @@ void interpolate(
         const LookUpFunc& lookup,
         idx_t* argmins) {
     std::unordered_map<idx_t, idx_t> idx_to_col;
-    for (idx_t idx = 0; idx < cols.size(); ++idx) {
+    for (size_t idx = 0; idx < cols.size(); ++idx) {
         idx_to_col[cols[idx]] = idx;
     }
 
     idx_t start = 0;
-    for (idx_t r = 0; r < rows.size(); r += 2) {
+    for (size_t r = 0; r < rows.size(); r += 2) {
         idx_t row = rows[r];
         idx_t end = cols.size() - 1;
         if (r < rows.size() - 1) {
@@ -107,7 +107,7 @@ void smawk_impl(
 
     // call recursively on odd-indexed rows
     std::vector<idx_t> odd_rows;
-    for (idx_t i = 1; i < rows.size(); i += 2) {
+    for (size_t i = 1; i < rows.size(); i += 2) {
         odd_rows.push_back(rows[i]);
     }
     smawk_impl(odd_rows, cols, lookup, argmins);
@@ -175,10 +175,10 @@ class Matrix {
     idx_t ncols;
 
    public:
-    Matrix(idx_t nrows, idx_t ncols) {
-        this->nrows = nrows;
-        this->ncols = ncols;
-        data.resize(nrows * ncols);
+    Matrix(idx_t nrows_in, idx_t ncols_in) {
+        this->nrows = nrows_in;
+        this->ncols = ncols_in;
+        data.resize(nrows_in * ncols_in);
     }
 
     inline T& at(idx_t i, idx_t j) {
@@ -240,14 +240,14 @@ double kmeans1d(const float* x, size_t n, size_t nclusters, float* centroids) {
     Matrix<float> D(nclusters, n);
     Matrix<idx_t> T(nclusters, n);
 
-    for (idx_t m = 0; m < n; m++) {
+    for (idx_t m = 0; m < static_cast<idx_t>(n); m++) {
         D.at(0, m) = CC(0, m);
         T.at(0, m) = 0;
     }
 
     std::vector<idx_t> indices(nclusters, 0);
 
-    for (idx_t k = 1; k < nclusters; ++k) {
+    for (idx_t k = 1; k < static_cast<idx_t>(nclusters); ++k) {
         // we define C here
         auto C = [&D, &CC, &k](idx_t m, idx_t i) {
             if (i == 0) {
@@ -259,7 +259,7 @@ double kmeans1d(const float* x, size_t n, size_t nclusters, float* centroids) {
 
         std::vector<idx_t> argmins(n); // argmin of each row
         smawk(n, n, C, argmins.data());
-        for (idx_t m = 0; m < argmins.size(); m++) {
+        for (size_t m = 0; m < argmins.size(); m++) {
             idx_t idx = argmins[m];
             D.at(k, m) = C(m, idx);
             T.at(k, m) = idx;

data/vendor/faiss/faiss/impl/lattice_Zn.cpp

@@ -18,6 +18,9 @@
 #include <queue>
 #include <unordered_set>
 
+#include <faiss/impl/FaissAssert.h>
+
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
 namespace faiss {
@@ -38,7 +41,7 @@ struct Comb {
     std::vector<uint64_t> tab; // Pascal's triangle
     int nmax;
 
-    explicit Comb(int nmax) : nmax(nmax) {
+    explicit Comb(int nmax_in) : nmax(nmax_in) {
         tab.resize(nmax * nmax, 0);
         tab[0] = 1;
         for (int i = 1; i < nmax; i++) {
@@ -169,9 +172,9 @@ void repeats_decode_64(
 
 } // anonymous namespace
 
-Repeats::Repeats(int dim, const float* c) : dim(dim) {
+Repeats::Repeats(int dim_in, const float* c) : dim(dim_in) {
     for (int i = 0; i < dim; i++) {
-        int j = 0;
+        size_t j = 0;
         for (;;) {
             if (j == repeats.size()) {
                 repeats.push_back(Repeat{c[i], 1});
@@ -189,7 +192,7 @@ Repeats::Repeats(int dim, const float* c) : dim(dim) {
 uint64_t Repeats::count() const {
     uint64_t accu = 1;
     int remain = dim;
-    for (int i = 0; i < repeats.size(); i++) {
+    for (size_t i = 0; i < repeats.size(); i++) {
         accu *= comb(remain, repeats[i].n);
         remain -= repeats[i].n;
     }
@@ -272,7 +275,7 @@ void EnumeratedVectors::encode_multi(size_t n, const float* c, uint64_t* codes)
 #pragma omp parallel if (n > 1000)
     {
 #pragma omp for
-        for (int i = 0; i < n; i++) {
+        for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
             codes[i] = encode(c + i * dim);
         }
     }
@@ -283,7 +286,7 @@ void EnumeratedVectors::decode_multi(size_t n, const uint64_t* codes, float* c)
 #pragma omp parallel if (n > 1000)
     {
 #pragma omp for
-        for (int i = 0; i < n; i++) {
+        for (int64_t i = 0; i < static_cast<int64_t>(n); i++) {
             decode(codes[i], c + i * dim);
         }
     }
@@ -302,25 +305,33 @@ void EnumeratedVectors::find_nn(
     }
 
     std::vector<float> c(dim);
-    for (size_t i = 0; i < nc; i++) {
-        uint64_t code = codes[nc];
-        decode(code, c.data());
-        for (size_t j = 0; j < nq; j++) {
-            const float* x = xq + j * dim;
-            float dis = fvec_inner_product(x, c.data(), dim);
-            if (dis > distances[j]) {
-                distances[j] = dis;
-                labels[j] = i;
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t i = 0; i < nc; i++) {
+            uint64_t code = codes[nc];
+            decode(code, c.data());
+            for (size_t j = 0; j < nq; j++) {
+                const float* x = xq + j * dim;
+                float dis = fvec_inner_product<SL>(x, c.data(), dim);
+                if (dis > distances[j]) {
+                    distances[j] = dis;
+                    labels[j] = i;
+                }
             }
         }
-    }
+    });
 }
 
 /**********************************************************
  * ZnSphereSearch
  **********************************************************/
 
-ZnSphereSearch::ZnSphereSearch(int dim, int r2) : dimS(dim), r2(r2) {
+ZnSphereSearch::ZnSphereSearch(int dim, int r2_in) : dimS(dim), r2(r2_in) {
+    FAISS_THROW_IF_NOT_MSG(
+            dim > 0 && dim <= 64, "ZnSphereSearch: dim must be in [1, 64]");
+    FAISS_THROW_IF_NOT_MSG(
+            r2_in >= 0 && r2_in <= 512,
+            "ZnSphereSearch: r2 must be in [0, 512] to avoid"
+            " excessive computation in sum_of_sq");
     voc = sum_of_sq(r2, int(ceil(sqrt(r2)) + 1), dim);
     natom = voc.size() / dim;
 }
@@ -355,13 +366,15 @@ float ZnSphereSearch::search(
     // find best
     int ibest = -1;
     float dpbest = -100;
-    for (int i = 0; i < natom; i++) {
-        float dp = fvec_inner_product(voc.data() + i * dim, xperm, dim);
-        if (dp > dpbest) {
-            dpbest = dp;
-            ibest = i;
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (int i = 0; i < natom; i++) {
+            float dp = fvec_inner_product<SL>(voc.data() + i * dim, xperm, dim);
+            if (dp > dpbest) {
+                dpbest = dp;
+                ibest = i;
+            }
         }
-    }
+    });
     // revert sort
     const float* cin = voc.data() + ibest * dim;
     for (int i = 0; i < dim; i++) {
@@ -391,8 +404,8 @@ void ZnSphereSearch::search_multi(
  * ZnSphereCodec
  **********************************************************/
 
-ZnSphereCodec::ZnSphereCodec(int dim, int r2)
-        : ZnSphereSearch(dim, r2), EnumeratedVectors(dim) {
+ZnSphereCodec::ZnSphereCodec(int dim_in, int r2_in)
+        : ZnSphereSearch(dim_in, r2_in), EnumeratedVectors(dim_in) {
     nv = 0;
     for (int i = 0; i < natom; i++) {
         Repeats repeats(dim, &voc[i * dim]);
@@ -484,16 +497,30 @@ void ZnSphereCodecRec::set_nv_cum(int ld, int r2t, int r2a, uint64_t cum) {
     all_nv_cum[(ld * (r2 + 1) + r2t) * (r2 + 1) + r2a] = cum;
 }
 
-ZnSphereCodecRec::ZnSphereCodecRec(int dim, int r2)
-        : EnumeratedVectors(dim), r2(r2) {
+ZnSphereCodecRec::ZnSphereCodecRec(int dim_in, int r2_in)
+        : EnumeratedVectors(dim_in), r2(r2_in) {
+    FAISS_THROW_IF_NOT_MSG(
+            dim_in > 0 && r2_in >= 0, "invalid ZnSphereCodecRec parameters");
     log2_dim = 0;
     while (dim > (1 << log2_dim)) {
         log2_dim++;
     }
-    assert(dim == (1 << log2_dim) || !"dimension must be a power of 2");
+    assert(dim == (1 << log2_dim) && "dimension must be a power of 2");
+
+    // Validate allocation sizes to avoid null pointer dereference on
+    // allocation failure. The cumulative table has O(r2^2) entries.
+    size_t nv_size = (size_t)(log2_dim + 1) * (r2 + 1);
+    size_t nv_cum_size = nv_size * (r2 + 1);
+    FAISS_THROW_IF_NOT_MSG(
+            nv_cum_size / (r2 + 1) == nv_size,
+            "ZnSphereCodecRec: allocation size overflow");
+    // Cap at ~1GB worth of uint64_t entries
+    FAISS_THROW_IF_NOT_MSG(
+            nv_cum_size <= (size_t(1) << 27),
+            "ZnSphereCodecRec: r2 too large, would require excessive memory");
 
-    all_nv.resize((log2_dim + 1) * (r2 + 1));
-    all_nv_cum.resize((log2_dim + 1) * (r2 + 1) * (r2 + 1));
+    all_nv.resize(nv_size);
+    all_nv_cum.resize(nv_cum_size);
 
     for (int r2a = 0; r2a <= r2; r2a++) {
         int r = int(sqrt(r2a));
@@ -506,13 +533,13 @@ ZnSphereCodecRec::ZnSphereCodecRec(int dim, int r2)
 
     for (int ld = 1; ld <= log2_dim; ld++) {
         for (int r2sub = 0; r2sub <= r2; r2sub++) {
-            uint64_t nv = 0;
+            uint64_t nv_acc = 0;
             for (int r2a = 0; r2a <= r2sub; r2a++) {
                 int r2b = r2sub - r2a;
-                set_nv_cum(ld, r2sub, r2a, nv);
-                nv += get_nv(ld - 1, r2a) * get_nv(ld - 1, r2b);
+                set_nv_cum(ld, r2sub, r2a, nv_acc);
+                nv_acc += get_nv(ld - 1, r2a) * get_nv(ld - 1, r2b);
             }
-            all_nv[ld * (r2 + 1) + r2sub] = nv;
+            all_nv[ld * (r2 + 1) + r2sub] = nv_acc;
         }
     }
     nv = get_nv(log2_dim, r2);
@@ -527,17 +554,31 @@ ZnSphereCodecRec::ZnSphereCodecRec(int dim, int r2)
     int cache_level = std::min(3, log2_dim - 1);
     decode_cache_ld = 0;
     assert(cache_level <= log2_dim);
+
     decode_cache.resize((r2 + 1));
 
+    // The decode cache stores total_cache_entries * dimsub floats.
+    // Cap at 2^27 total floats (~512 MB), aligned with the nv_cum
+    // memory cap above which also uses 2^27 entries.  The entry count
+    // grows as the number of lattice points in dimension
+    // 2^cache_level, which is O(r2^(dim/2)) -- much faster than the
+    // O(r2^2) growth of nv_cum.
+    size_t total_cache_entries = 0;
+    int dimsub = (1 << cache_level);
+
     for (int r2sub = 0; r2sub <= r2; r2sub++) {
         int ld = cache_level;
         uint64_t nvi = get_nv(ld, r2sub);
+        total_cache_entries += nvi;
+        FAISS_THROW_IF_NOT_MSG(
+                total_cache_entries <= (size_t(1) << 27) / dimsub,
+                "ZnSphereCodecRec: r2 too large, decode cache "
+                "would require excessive memory");
         std::vector<float>& cache = decode_cache[r2sub];
-        int dimsub = (1 << cache_level);
         cache.resize(nvi * dimsub);
         std::vector<float> c(dim);
         uint64_t code0 = get_nv_cum(cache_level + 1, r2, r2 - r2sub);
-        for (int i = 0; i < nvi; i++) {
+        for (uint64_t i = 0; i < nvi; i++) {
             decode(i + code0, c.data());
             memcpy(&cache[i * dimsub],
                    c.data() + dim - dimsub,
@@ -643,10 +684,10 @@ void ZnSphereCodecRec::decode(uint64_t code, float* c) const {
 }
 
 // if not use_rec, instantiate an arbitrary harmless znc_rec
-ZnSphereCodecAlt::ZnSphereCodecAlt(int dim, int r2)
-        : ZnSphereCodec(dim, r2),
-          use_rec((dim & (dim - 1)) == 0),
-          znc_rec(use_rec ? dim : 8, use_rec ? r2 : 14) {}
+ZnSphereCodecAlt::ZnSphereCodecAlt(int dim_in, int r2_in)
+        : ZnSphereCodec(dim_in, r2_in),
+          use_rec((dim_in & (dim_in - 1)) == 0),
+          znc_rec(use_rec ? dim_in : 8, use_rec ? r2_in : 14) {}
 
 uint64_t ZnSphereCodecAlt::encode(const float* x) const {
     if (!use_rec) {