refactor memory management

Signed-off-by: Alexandr Guzhva <[email protected]>
zilliztech · May 15, 2024 · 52fdec9 · 52fdec9
1 parent fe5a052
commit 52fdec9
Show file tree

Hide file tree

Showing 10 changed files with 155 additions and 98 deletions.
diff --git a/include/knowhere/dataset.h b/include/knowhere/dataset.h
@@ -22,6 +22,7 @@
 
 #include "comp/index_param.h"
 #include "knowhere/sparse_utils.h"
+#include "knowhere/range_util.h"
 
 namespace knowhere {
 
@@ -71,18 +72,36 @@ class DataSet : public std::enable_shared_from_this<const DataSet> {
         this->data_[meta::DISTANCE] = Var(std::in_place_index<0>, dis);
     }
 
+    void
+    SetDistance(std::unique_ptr<float[]>&& dis) {
+        std::unique_lock lock(mutex_);
+        this->data_[meta::DISTANCE] = Var(std::in_place_index<0>, dis.release());
+    }
+
     void
     SetLims(const size_t* lims) {
         std::unique_lock lock(mutex_);
         this->data_[meta::LIMS] = Var(std::in_place_index<1>, lims);
     }
 
+    void
+    SetLims(std::unique_ptr<size_t[]>&& lims) {
+        std::unique_lock lock(mutex_);
+        this->data_[meta::LIMS] = Var(std::in_place_index<1>, lims.release());
+    }
+
     void
     SetIds(const int64_t* ids) {
         std::unique_lock lock(mutex_);
         this->data_[meta::IDS] = Var(std::in_place_index<2>, ids);
     }
 
+    void
+    SetIds(std::unique_ptr<int64_t[]>&& ids) {
+        std::unique_lock lock(mutex_);
+        this->data_[meta::IDS] = Var(std::in_place_index<2>, ids.release());
+    }
+
     /**
      * For dense float vector, tensor is a rows * dim float array
      * For sparse float vector, tensor is pointer to sparse::Sparse<float>*
@@ -94,6 +113,18 @@ class DataSet : public std::enable_shared_from_this<const DataSet> {
         this->data_[meta::TENSOR] = Var(std::in_place_index<3>, tensor);
     }
 
+    void
+    SetTensor(std::unique_ptr<uint8_t[]>&& tensor) {
+        std::unique_lock lock(mutex_);
+        this->data_[meta::TENSOR] = Var(std::in_place_index<3>, tensor.release());
+    }
+
+    void
+    SetTensor(std::unique_ptr<float[]>&& tensor) {
+        std::unique_lock lock(mutex_);
+        this->data_[meta::TENSOR] = Var(std::in_place_index<3>, tensor.release());
+    }
+
     void
     SetRows(const int64_t rows) {
         std::unique_lock lock(mutex_);
@@ -284,6 +315,26 @@ GenResultDataSet(const int64_t rows, const int64_t dim, const void* tensor) {
     return ret_ds;
 }
 
+inline DataSetPtr
+GenResultDataSet(const int64_t rows, const int64_t dim, std::unique_ptr<uint8_t[]>&& tensor) {
+    auto ret_ds = std::make_shared<DataSet>();
+    ret_ds->SetRows(rows);
+    ret_ds->SetDim(dim);
+    ret_ds->SetTensor(std::move(tensor));
+    ret_ds->SetIsOwner(true);
+    return ret_ds;
+}
+
+inline DataSetPtr
+GenResultDataSet(const int64_t rows, const int64_t dim, std::unique_ptr<float[]>&& tensor) {
+    auto ret_ds = std::make_shared<DataSet>();
+    ret_ds->SetRows(rows);
+    ret_ds->SetDim(dim);
+    ret_ds->SetTensor(std::move(tensor));
+    ret_ds->SetIsOwner(true);
+    return ret_ds;
+}
+
 inline DataSetPtr
 #ifdef NOT_COMPILE_FOR_SWIG
 GenResultDataSet(const int64_t nq, const int64_t topk, const int64_t* ids, const float* distance) {
@@ -299,6 +350,22 @@ GenResultDataSet(const int64_t nq, const int64_t topk, const void* ids, const fl
     return ret_ds;
 }
 
+inline DataSetPtr
+GenResultDataSet(
+    const int64_t nq,
+    const int64_t topk,
+    std::unique_ptr<int64_t[]>&& ids,
+    std::unique_ptr<float[]>&& distance
+) {
+    auto ret_ds = std::make_shared<DataSet>();
+    ret_ds->SetRows(nq);
+    ret_ds->SetDim(topk);
+    ret_ds->SetIds(std::move(ids));
+    ret_ds->SetDistance(std::move(distance));
+    ret_ds->SetIsOwner(true);
+    return ret_ds;
+}
+
 inline DataSetPtr
 #ifdef NOT_COMPILE_FOR_SWIG
 GenResultDataSet(const int64_t nq, const int64_t* ids, const float* distance, const size_t* lims) {
@@ -314,6 +381,17 @@ GenResultDataSet(const int64_t nq, const void* ids, const float* distance, const
     return ret_ds;
 }
 
+inline DataSetPtr
+GenResultDataSet(const int64_t nq, RangeSearchResult&& range_search_result) {
+    auto ret_ds = std::make_shared<DataSet>();
+    ret_ds->SetRows(nq);
+    ret_ds->SetIds(std::move(range_search_result.labels));
+    ret_ds->SetDistance(std::move(range_search_result.distances));
+    ret_ds->SetLims(std::move(range_search_result.lims));
+    ret_ds->SetIsOwner(true);
+    return ret_ds;
+}
+
 inline DataSetPtr
 GenResultDataSet(const std::string& json_info, const std::string& json_id_set) {
     auto ret_ds = std::make_shared<DataSet>();

diff --git a/include/knowhere/index/index_node.h b/include/knowhere/index/index_node.h
@@ -134,12 +134,8 @@ class IndexNode : public Object {
         }
 #endif
 
-        int64_t* ids = nullptr;
-        float* dis = nullptr;
-        size_t* lims = nullptr;
-        GetRangeSearchResult(result_dist_array, result_id_array, similarity_metric, nq, radius, range_filter, dis, ids,
-                             lims);
-        return GenResultDataSet(nq, ids, dis, lims);
+        auto range_search_result = GetRangeSearchResult(result_dist_array, result_id_array, similarity_metric, nq, radius, range_filter);
+        return GenResultDataSet(nq, std::move(range_search_result));
     }
 
     virtual expected<DataSetPtr>

diff --git a/include/knowhere/range_util.h b/include/knowhere/range_util.h
@@ -11,6 +11,8 @@
 
 #pragma once
 
+#include <cstddef>
+#include <memory>
 #include <vector>
 
 #include "knowhere/bitsetview.h"
@@ -26,9 +28,15 @@ void
 FilterRangeSearchResultForOneNq(std::vector<float>& distances, std::vector<int64_t>& labels, const bool is_ip,
                                 const float radius, const float range_filter);
 
-void
+struct RangeSearchResult {
+    std::unique_ptr<float[]> distances;
+    std::unique_ptr<int64_t[]> labels;
+    std::unique_ptr<size_t[]> lims;
+};
+
+RangeSearchResult
 GetRangeSearchResult(const std::vector<std::vector<float>>& result_distances,
                      const std::vector<std::vector<int64_t>>& result_labels, const bool is_ip, const int64_t nq,
-                     const float radius, const float range_filter, float*& distances, int64_t*& labels, size_t*& lims);
+                     const float radius, const float range_filter);
 
 }  // namespace knowhere
diff --git a/src/common/comp/brute_force.cc b/src/common/comp/brute_force.cc
@@ -153,7 +153,7 @@ BruteForce::Search(const DataSetPtr base_dataset, const DataSetPtr query_dataset
     if (ret != Status::success) {
         return expected<DataSetPtr>::Err(ret, "failed to brute force search");
     }
-    auto res = GenResultDataSet(nq, cfg.k.value(), labels.release(), distances.release());
+    auto res = GenResultDataSet(nq, cfg.k.value(), std::move(labels), std::move(distances));
 
 #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT)
     if (cfg.trace_id.has_value()) {
@@ -432,11 +432,8 @@ BruteForce::RangeSearch(const DataSetPtr base_dataset, const DataSetPtr query_da
         return expected<DataSetPtr>::Err(ret, "failed to brute force search");
     }
 
-    int64_t* ids = nullptr;
-    float* distances = nullptr;
-    size_t* lims = nullptr;
-    GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius, range_filter, distances, ids, lims);
-    auto res = GenResultDataSet(nq, ids, distances, lims);
+    auto range_search_result = GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius, range_filter);
+    auto res = GenResultDataSet(nq, std::move(range_search_result));
 
 #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT)
     if (cfg.trace_id.has_value()) {
@@ -550,7 +547,7 @@ BruteForce::SearchSparse(const DataSetPtr base_dataset, const DataSetPtr query_d
     auto distances = std::make_unique<float[]>(nq * topk);
 
     SearchSparseWithBuf(base_dataset, query_dataset, labels.get(), distances.get(), config, bitset);
-    return GenResultDataSet(nq, topk, labels.release(), distances.release());
+    return GenResultDataSet(nq, topk, std::move(labels), std::move(distances));
 }
 
 template <typename DataType>

diff --git a/src/common/range_util.cc b/src/common/range_util.cc
@@ -12,7 +12,11 @@
 #include "knowhere/range_util.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <cinttypes>
+#include <memory>
+#include <tuple>
+#include <vector>
 
 #include "knowhere/log.h"
 namespace knowhere {
@@ -41,16 +45,16 @@ FilterRangeSearchResultForOneNq(std::vector<float>& distances, std::vector<int64
     }
 }
 
-void
+RangeSearchResult
 GetRangeSearchResult(const std::vector<std::vector<float>>& result_distances,
                      const std::vector<std::vector<int64_t>>& result_labels, const bool is_ip, const int64_t nq,
-                     const float radius, const float range_filter, float*& distances, int64_t*& labels, size_t*& lims) {
+                     const float radius, const float range_filter) {
     KNOWHERE_THROW_IF_NOT_FMT(result_distances.size() == (size_t)nq, "result distances size %ld not equal to %" SCNd64,
                               result_distances.size(), nq);
     KNOWHERE_THROW_IF_NOT_FMT(result_labels.size() == (size_t)nq, "result labels size %ld not equal to %" SCNd64,
                               result_labels.size(), nq);
 
-    lims = new size_t[nq + 1];
+    auto lims = std::make_unique<size_t[]>(nq + 1);
     lims[0] = 0;
     // all distances must be in range scope
     for (int64_t i = 0; i < nq; i++) {
@@ -61,13 +65,19 @@ GetRangeSearchResult(const std::vector<std::vector<float>>& result_distances,
     LOG_KNOWHERE_DEBUG_ << "Range search: is_ip " << (is_ip ? "True" : "False") << ", radius " << radius
                         << ", range_filter " << range_filter << ", total result num " << total_valid;
 
-    distances = new float[total_valid];
-    labels = new int64_t[total_valid];
+    auto distances = std::make_unique<float[]>(total_valid);
+    auto labels = std::make_unique<int64_t[]>(total_valid);
 
     for (auto i = 0; i < nq; i++) {
-        std::copy_n(result_distances[i].data(), lims[i + 1] - lims[i], distances + lims[i]);
-        std::copy_n(result_labels[i].data(), lims[i + 1] - lims[i], labels + lims[i]);
+        std::copy_n(result_distances[i].data(), lims[i + 1] - lims[i], distances.get() + lims[i]);
+        std::copy_n(result_labels[i].data(), lims[i + 1] - lims[i], labels.get() + lims[i]);
     }
+
+    return RangeSearchResult{
+        .distances = std::move(distances),
+        .labels = std::move(labels),
+        .lims = std::move(lims)
+    };
 }
 
 }  // namespace knowhere
diff --git a/src/index/diskann/diskann.cc b/src/index/diskann/diskann.cc
@@ -557,7 +557,7 @@ DiskANNIndexNode<DataType>::Search(const DataSet& dataset, const Config& cfg, co
         return expected<DataSetPtr>::Err(Status::diskann_inner_error, "some search failed");
     }
 
-    auto res = GenResultDataSet(nq, k, p_id.release(), p_dist.release());
+    auto res = GenResultDataSet(nq, k, p_id, p_dist);
 
     // set visit_info json string into result dataset
     if (feder_result != nullptr) {
@@ -600,10 +600,6 @@ DiskANNIndexNode<DataType>::RangeSearch(const DataSet& dataset, const Config& cf
     auto nq = dataset.GetRows();
     auto xq = static_cast<const DataType*>(dataset.GetTensor());
 
-    int64_t* p_id = nullptr;
-    DistType* p_dist = nullptr;
-    size_t* p_lims = nullptr;
-
     std::vector<std::vector<int64_t>> result_id_array(nq);
     std::vector<std::vector<DistType>> result_dist_array(nq);
 
@@ -628,9 +624,8 @@ DiskANNIndexNode<DataType>::RangeSearch(const DataSet& dataset, const Config& cf
         return expected<DataSetPtr>::Err(Status::diskann_inner_error, "some search failed");
     }
 
-    GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius, search_conf.range_filter.value(),
-                         p_dist, p_id, p_lims);
-    return GenResultDataSet(nq, p_id, p_dist, p_lims);
+    auto range_search_result = GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius, search_conf.range_filter.value());
+    return GenResultDataSet(nq, std::move(range_search_result));
 }
 
 /*

diff --git a/src/index/flat/flat.cc b/src/index/flat/flat.cc
@@ -156,9 +156,7 @@ class FlatIndexNode : public IndexNode {
         float range_filter = f_cfg.range_filter.value();
         bool is_ip = (index_->metric_type == faiss::METRIC_INNER_PRODUCT);
 
-        int64_t* ids = nullptr;
-        float* distances = nullptr;
-        size_t* lims = nullptr;
+        RangeSearchResult range_search_result;
 
         std::vector<std::vector<int64_t>> result_id_array(nq);
         std::vector<std::vector<float>> result_dist_array(nq);
@@ -208,14 +206,13 @@ class FlatIndexNode : public IndexNode {
             }
             // wait for the completion
             WaitAllSuccess(futs);
-            GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius, range_filter, distances, ids,
-                                 lims);
+            range_search_result = GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius, range_filter);
         } catch (const std::exception& e) {
             LOG_KNOWHERE_WARNING_ << "error inner faiss: " << e.what();
             return expected<DataSetPtr>::Err(Status::faiss_inner_error, e.what());
         }
 
-        return GenResultDataSet(nq, ids, distances, lims);
+        return GenResultDataSet(nq, std::move(range_search_result));
     }
 
     expected<DataSetPtr>

diff --git a/src/index/hnsw/hnsw.cc b/src/index/hnsw/hnsw.cc
@@ -217,7 +217,7 @@ class HnswIndexNode : public IndexNode {
         }
         WaitAllSuccess(futs);
 
-        auto res = GenResultDataSet(nq, k, p_id.release(), p_dist.release());
+        auto res = GenResultDataSet(nq, k, std::move(p_id), std::move(p_dist));
 
         // set visit_info json string into result dataset
         if (feder_result != nullptr) {
@@ -358,15 +358,10 @@ class HnswIndexNode : public IndexNode {
         }
         WaitAllSuccess(futs);
 
-        int64_t* ids = nullptr;
-        DistType* dis = nullptr;
-        size_t* lims = nullptr;
-
         // filter range search result
-        GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius_for_filter, range_filter, dis, ids,
-                             lims);
+        auto range_search_result = GetRangeSearchResult(result_dist_array, result_id_array, is_ip, nq, radius_for_filter, range_filter);
 
-        auto res = GenResultDataSet(nq, ids, dis, lims);
+        auto res = GenResultDataSet(nq, std::move(range_search_result));
 
         // set visit_info json string into result dataset
         if (feder_result != nullptr) {