Add HierarchicalNSW::indexFileSize() function for precise memory foot…

…print control
nmslib · Jan 18, 2023 · 208a0d7 · 208a0d7
1 parent 2c6f244
commit 208a0d7
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/hnswlib/hnswalg.h b/hnswlib/hnswalg.h
@@ -595,6 +595,32 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         max_elements_ = new_max_elements;
     }
 
+    size_t indexFileSize() const {
+        size_t size = 0;
+        size += sizeof(offsetLevel0_);
+        size += sizeof(max_elements_);
+        size += sizeof(cur_element_count);
+        size += sizeof(size_data_per_element_);
+        size += sizeof(label_offset_);
+        size += sizeof(offsetData_);
+        size += sizeof(maxlevel_);
+        size += sizeof(enterpoint_node_);
+        size += sizeof(maxM_);
+
+        size += sizeof(maxM0_);
+        size += sizeof(M_);
+        size += sizeof(mult_);
+        size += sizeof(ef_construction_);
+
+        size += cur_element_count * size_data_per_element_;
+
+        for (size_t i = 0; i < cur_element_count; i++) {
+            unsigned int linkListSize = element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0;
+            size += sizeof(linkListSize);
+            size += linkListSize;
+        }
+        return size;
+    }
 
     void saveIndex(const std::string &location) {
         std::ofstream output(location, std::ios::binary);

diff --git a/python_bindings/bindings.cpp b/python_bindings/bindings.cpp
@@ -218,6 +218,9 @@ class Index {
         this->num_threads_default = num_threads;
     }
 
+    void indexFileSize() const {
+        appr_alg->indexFileSize();
+    }
 
     void saveIndex(const std::string &path_to_index) {
         appr_alg->saveIndex(path_to_index);
@@ -904,6 +907,7 @@ PYBIND11_PLUGIN(hnswlib) {
         .def("get_ids_list", &Index<float>::getIdsList)
         .def("set_ef", &Index<float>::set_ef, py::arg("ef"))
         .def("set_num_threads", &Index<float>::set_num_threads, py::arg("num_threads"))
+        .def("index_file_size", &Index<float>::indexFileSize)
         .def("save_index", &Index<float>::saveIndex, py::arg("path_to_index"))
         .def("load_index",
             &Index<float>::loadIndex,