Declare saveIndex/loadIndex as templates to use custom streams

drons · Apr 13, 2024 · d0e9347 · d0e9347
1 parent e2ca570
commit d0e9347
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/hnswlib/hnswalg.h b/hnswlib/hnswalg.h
@@ -684,7 +684,11 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
     void saveIndex(const std::string &location) {
         std::ofstream output(location, std::ios::binary);
-
+        saveIndex(output);
+    }
+
+    template<class stream_t>
+    void saveIndex(stream_t &output) {
         writeBinaryPOD(output, offsetLevel0_);
         writeBinaryPOD(output, max_elements_);
         writeBinaryPOD(output, cur_element_count);
@@ -711,17 +715,20 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
         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);
 
         if (!input.is_open())
             throw std::runtime_error("Cannot open file");
+        loadIndex(input, s, max_elements_i);
+    }
 
+    template<class stream_t>
+    void loadIndex(stream_t &input, SpaceInterface<dist_t> *s, size_t max_elements_i = 0) {
         clear();
         // get file size:
         input.seekg(0, input.end);
-        std::streampos total_filesize = input.tellg();
+        auto total_filesize = input.tellg();
         input.seekg(0, input.beg);
 
         readBinaryPOD(input, offsetLevel0_);