Added generic serial / deserialization interfaces

README.md

@@ -3,6 +3,10 @@ Header-only C++ HNSW implementation with python bindings, insertions and updates
 
 **NEWS:**
 
+**version 0.8.1**
+
+* Added generic serialization / deserialization interfaces for HierarchicalNSW and BruteforceSearch that take `std::ostream` / `std::istream` arguments
+
 **version 0.8.0** 
 
 * Multi-vector document search and epsilon search (for now, only in C++)

hnswlib/bruteforce.h

@@ -4,6 +4,7 @@
 #include <mutex>
 #include <algorithm>
 #include <assert.h>
+#include <iostream>
 
 namespace hnswlib {
 template<typename dist_t>
@@ -135,24 +136,25 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
     }
 
 
-    void saveIndex(const std::string &location) {
-        std::ofstream output(location, std::ios::binary);
-        std::streampos position;
-
+    void saveIndex(std::ostream &output) {
         writeBinaryPOD(output, maxelements_);
         writeBinaryPOD(output, size_per_element_);
         writeBinaryPOD(output, cur_element_count);
 
         output.write(data_, maxelements_ * size_per_element_);
+    }
+
+
+    void saveIndex(const std::string &location) {
+        std::ofstream output(location, std::ios::binary);
+
+        saveIndex(output);
 
         output.close();
     }
 
 
-    void loadIndex(const std::string &location, SpaceInterface<dist_t> *s) {
-        std::ifstream input(location, std::ios::binary);
-        std::streampos position;
-
+    void loadIndex(std::istream &input, SpaceInterface<dist_t> *s) {
         readBinaryPOD(input, maxelements_);
         readBinaryPOD(input, size_per_element_);
         readBinaryPOD(input, cur_element_count);
@@ -166,6 +168,13 @@ class BruteforceSearch : public AlgorithmInterface<dist_t> {
             throw std::runtime_error("Not enough memory: loadIndex failed to allocate data");
 
         input.read(data_, maxelements_ * size_per_element_);
+    }
+
+
+    void loadIndex(const std::string &location, SpaceInterface<dist_t> *s) {
+        std::ifstream input(location, std::ios::binary);
+
+        loadIndex(input, s);
 
         input.close();
     }

hnswlib/hnswalg.h

@@ -9,6 +9,7 @@
 #include <unordered_set>
 #include <list>
 #include <memory>
+#include <iostream>
 
 namespace hnswlib {
 typedef unsigned int tableint;
@@ -682,10 +683,8 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         return size;
     }
 
-    void saveIndex(const std::string &location) {
-        std::ofstream output(location, std::ios::binary);
-        std::streampos position;
 
+    void saveIndex(std::ostream &output) {
         writeBinaryPOD(output, offsetLevel0_);
         writeBinaryPOD(output, max_elements_);
         writeBinaryPOD(output, cur_element_count);
@@ -709,16 +708,17 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
             if (linkListSize)
                 output.write(linkLists_[i], linkListSize);
         }
-        output.close();
     }
 
 
-    void loadIndex(const std::string &location, SpaceInterface<dist_t> *s, size_t max_elements_i = 0) {
-        std::ifstream input(location, std::ios::binary);
+    void saveIndex(const std::string &location) {
+        std::ofstream output(location, std::ios::binary);
+        saveIndex(output);
+        output.close();
+    }
 
-        if (!input.is_open())
-            throw std::runtime_error("Cannot open file");
 
+    void loadIndex(std::istream &input, SpaceInterface<dist_t> *s, size_t max_elements_i = 0) {
         clear();
         // get file size:
         input.seekg(0, input.end);
@@ -815,6 +815,16 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
                 if (allow_replace_deleted_) deleted_elements.insert(i);
             }
         }
+    }
+
+
+    void loadIndex(const std::string &location, SpaceInterface<dist_t> *s, size_t max_elements_i = 0) {
+        std::ifstream input(location, std::ios::binary);
+
+        if (!input.is_open())
+            throw std::runtime_error("Cannot open file");
+
+        loadIndex(input, s, max_elements_i);
 
         input.close();
 

setup.py

@@ -8,7 +8,7 @@
 from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
-__version__ = '0.8.0'
+__version__ = '0.8.1'
 
 
 include_dirs = [