From a9de0605f18f5b6a5e9cc887d946e04ef115a4a7 Mon Sep 17 00:00:00 2001
From: drons <drons@list.ru>
Date: Mon, 5 Dec 2022 20:00:13 +0300
Subject: [PATCH] Add HierarchicalNSW::indexFileSize() function for precise
 memory footprint control

---
 hnswlib/hnswalg.h            | 26 ++++++++++++++++++++++++++
 python_bindings/bindings.cpp |  4 ++++
 2 files changed, 30 insertions(+)

diff --git a/hnswlib/hnswalg.h b/hnswlib/hnswalg.h
index bef00170..6a5495c0 100644
--- 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
index 5153bb58..5391419c 100644
--- a/python_bindings/bindings.cpp
+++ b/python_bindings/bindings.cpp
@@ -218,6 +218,9 @@ class Index {
         this->num_threads_default = num_threads;
     }
 
+    size_t indexFileSize() const {
+        return 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,