Support deserialize from file and mmap (#722)

Signed-off-by: yah01 <yang.cen@zilliz.com>
zilliztech · Mar 27, 2023 · 7e52ec0 · 7e52ec0
1 parent bf2a600
commit 7e52ec0
Show file tree

Hide file tree

Showing 15 changed files with 185 additions and 225 deletions.
diff --git a/include/knowhere/config.h b/include/knowhere/config.h
@@ -418,6 +418,11 @@ class BaseConfig : public Config {
     }
 };
 
+struct LoadConfig {
+    // load with mmap
+    bool enable_mmap{false};
+};
+
 }  // namespace knowhere
 
 #endif /* CONFIG_H */
diff --git a/include/knowhere/index.h b/include/knowhere/index.h
@@ -12,6 +12,7 @@
 #ifndef INDEX_H
 #define INDEX_H
 
+#include "knowhere/config.h"
 #include "knowhere/index_node.h"
 #include "knowhere/log.h"
 
@@ -248,6 +249,11 @@ class Index {
         return this->node->Deserialize(binset);
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config = {}) {
+        return this->node->Deserialize(filename, config);
+    }
+
     int64_t
     Dim() const {
         return this->node->Dim();

diff --git a/include/knowhere/index_node.h b/include/knowhere/index_node.h
@@ -50,6 +50,9 @@ class IndexNode : public Object {
     virtual Status
     Deserialize(const BinarySet& binset) = 0;
 
+    virtual Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) = 0;
+
     virtual std::unique_ptr<BaseConfig>
     CreateConfig() const = 0;
 

diff --git a/include/knowhere/index_node_thread_pool_wrapper.h b/include/knowhere/index_node_thread_pool_wrapper.h
@@ -72,6 +72,11 @@ class IndexNodeThreadPoolWrapper : public IndexNode {
         return index_node_->Deserialize(binset);
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) {
+        return index_node_->DeserializeFromFile(filename, config);
+    }
+
     std::unique_ptr<BaseConfig>
     CreateConfig() const {
         return index_node_->CreateConfig();

diff --git a/src/index/annoy/annoy.cc b/src/index/annoy/annoy.cc
@@ -13,6 +13,7 @@
 #include "index/annoy/annoy_config.h"
 #include "kissrandom.h"
 #include "knowhere/comp/thread_pool.h"
+#include "knowhere/expected.h"
 #include "knowhere/factory.h"
 namespace knowhere {
 
@@ -218,6 +219,11 @@ class AnnoyIndexNode : public IndexNode {
         return Status::success;
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override {
+        return Status::not_implemented;
+    }
+
     std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         return std::make_unique<AnnoyConfig>();

diff --git a/src/index/diskann/diskann.cc b/src/index/diskann/diskann.cc
@@ -80,6 +80,11 @@ class DiskANNIndexNode : public IndexNode {
         return Status::not_implemented;
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override {
+        return Status::not_implemented;
+    }
+
     std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         return std::make_unique<DiskANNConfig>();

diff --git a/src/index/flat/flat.cc b/src/index/flat/flat.cc
@@ -271,6 +271,11 @@ class FlatIndexNode : public IndexNode {
         return Status::success;
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override {
+        return Status::not_implemented;
+    }
+
     std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         return std::make_unique<FlatConfig>();

diff --git a/src/index/flat_gpu/flat_gpu.cc b/src/index/flat_gpu/flat_gpu.cc
@@ -159,6 +159,11 @@ class GpuFlatIndexNode : public IndexNode {
         return Status::success;
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override {
+        return Status::not_implemented;
+    }
+
     virtual std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         return std::make_unique<GpuFlatConfig>();

diff --git a/src/index/hnsw/hnsw.cc b/src/index/hnsw/hnsw.cc
@@ -13,13 +13,19 @@
 
 #include <omp.h>
 
+#include <exception>
+#include <new>
+
 #include "common/range_util.h"
 #include "hnswlib/hnswalg.h"
 #include "hnswlib/hnswlib.h"
 #include "index/hnsw/hnsw_config.h"
 #include "knowhere/comp/index_param.h"
 #include "knowhere/comp/thread_pool.h"
+#include "knowhere/config.h"
+#include "knowhere/expected.h"
 #include "knowhere/factory.h"
+#include "knowhere/log.h"
 
 namespace knowhere {
 class HnswIndexNode : public IndexNode {
@@ -321,6 +327,22 @@ class HnswIndexNode : public IndexNode {
         return Status::success;
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override {
+        if (index_) {
+            delete index_;
+        }
+        try {
+            hnswlib::SpaceInterface<float>* space = nullptr;
+            index_ = new (std::nothrow) hnswlib::HierarchicalNSW<float>(space);
+            index_->loadIndex(filename, config);
+        } catch (std::exception& e) {
+            LOG_KNOWHERE_WARNING_ << "hnsw inner error, " << e.what();
+            return Status::hnsw_inner_error;
+        }
+        return Status::success;
+    }
+
     std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         return std::make_unique<HnswConfig>();

diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc
@@ -68,6 +68,8 @@ class IvfIndexNode : public IndexNode {
     Serialize(BinarySet& binset) const override;
     Status
     Deserialize(const BinarySet& binset) override;
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override;
     std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         if constexpr (std::is_same<faiss::IndexIVFFlat, T>::value) {
@@ -605,6 +607,12 @@ IvfIndexNode<T>::Deserialize(const BinarySet& binset) {
     return Status::success;
 }
 
+template <typename T>
+Status
+IvfIndexNode<T>::DeserializeFromFile(const std::string& filename, const LoadConfig& config) {
+    return Status::not_implemented;
+}
+
 template <>
 Status
 IvfIndexNode<faiss::IndexIVFFlat>::Deserialize(const BinarySet& binset) {

diff --git a/src/index/ivf_gpu/ivf_gpu.cc b/src/index/ivf_gpu/ivf_gpu.cc
@@ -226,6 +226,11 @@ class GpuIvfIndexNode : public IndexNode {
         return Status::success;
     }
 
+    Status
+    DeserializeFromFile(const std::string& filename, const LoadConfig& config) override {
+        return Status::not_implemented;
+    }
+
     virtual std::unique_ptr<BaseConfig>
     CreateConfig() const override {
         return std::make_unique<typename KnowhereConfigType<T>::Type>();

diff --git a/src/io/fileIO.h b/src/io/fileIO.h
@@ -0,0 +1,62 @@
+// Copyright (C) 2019-2023 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License
+
+#pragma once
+#include <sys/fcntl.h>
+#include <unistd.h>
+
+#include <stdexcept>
+
+namespace knowhere {
+struct FileReader {
+    int fd;
+    size_t size;
+
+    FileReader(const std::string& filename, bool auto_remove = false) {
+        fd = open(filename.data(), O_RDONLY);
+        if (fd < 0) {
+            std::runtime_error("Cannot open file");
+        }
+
+        size = lseek(fd, 0, SEEK_END);
+        lseek(fd, 0, SEEK_SET);
+
+        if (auto_remove) {
+            unlink(filename.data());
+        }
+    }
+
+    ssize_t
+    read(char* dst, size_t n) {
+        return ::read(fd, dst, n);
+    }
+
+    off_t
+    seek(off_t offset) {
+        return lseek(fd, offset, SEEK_SET);
+    }
+
+    off_t
+    advance(off_t offset) {
+        return lseek(fd, offset, SEEK_CUR);
+    }
+
+    off_t
+    offset() {
+        return lseek(fd, 0, SEEK_CUR);
+    }
+
+    int
+    close() {
+        return ::close(fd);
+    }
+};
+}  // namespace knowhere
diff --git a/thirdparty/hnswlib/hnswlib/bruteforce.h b/thirdparty/hnswlib/hnswlib/bruteforce.h