Delete RAW_DATA copy when load IVF_FLAT index data (#20274)

Signed-off-by: xige-16 <xi.ge@zilliz.com>

Signed-off-by: xige-16 <xi.ge@zilliz.com>

internal/core/src/index/CMakeLists.txt

@@ -14,6 +14,7 @@ set(INDEX_FILES
         Utils.cpp
         VectorMemIndex.cpp
         IndexFactory.cpp
+        VectorMemNMIndex.cpp
         )
 
 if ( BUILD_DISK_ANN STREQUAL "ON" )

internal/core/src/index/IndexFactory.cpp

@@ -16,6 +16,7 @@
 
 #include "index/IndexFactory.h"
 #include "index/VectorMemIndex.h"
+#include "index/VectorMemNMIndex.h"
 #include "index/Utils.h"
 #include "index/Meta.h"
 
@@ -87,6 +88,9 @@ IndexFactory::CreateVectorIndex(const CreateIndexInfo& create_index_info, storag
     }
 #endif
 
+    if (is_in_nm_list(index_type)) {
+        return std::make_unique<VectorMemNMIndex>(index_type, metric_type, index_mode);
+    }
     // create mem index
     return std::make_unique<VectorMemIndex>(index_type, metric_type, index_mode);
 }

internal/core/src/index/VectorMemIndex.cpp

@@ -43,17 +43,6 @@ VectorMemIndex::Serialize(const Config& config) {
     parse_config(serialize_config);
 
     auto ret = index_->Serialize(serialize_config);
-    auto index_type = GetIndexType();
-
-    if (is_in_nm_list(index_type)) {
-        auto deleter = [&](uint8_t*) {};  // avoid repeated deconstruction
-        auto raw_data = std::shared_ptr<uint8_t[]>(static_cast<uint8_t*>(raw_data_.data()), deleter);
-
-        //        std::shared_ptr<uint8_t[]> raw_data(new uint8_t[raw_data_.size()], std::default_delete<uint8_t[]>());
-        //        memcpy(raw_data.get(), raw_data_.data(), raw_data_.size());
-        ret.Append(RAW_DATA, raw_data, raw_data_.size());
-        // Disassemble will only divide the raw vectors, other keys were already divided
-    }
     milvus::Disassemble(ret);
 
     return ret;
@@ -62,19 +51,7 @@ VectorMemIndex::Serialize(const Config& config) {
 void
 VectorMemIndex::Load(const BinarySet& binary_set, const Config& config) {
     milvus::Assemble(const_cast<BinarySet&>(binary_set));
-
     index_->Load(binary_set);
-    auto& map_ = binary_set.binary_map_;
-    // copy RAW_DATA after index->Load(), since assemble is performed inside.
-    for (auto it = map_.begin(); it != map_.end(); ++it) {
-        if (it->first == RAW_DATA) {
-            raw_data_.clear();
-            auto data_size = it->second->size;
-            raw_data_.resize(data_size);
-            memcpy(raw_data_.data(), it->second->data.get(), data_size);
-            break;
-        }
-    }
     SetDim(index_->Dim());
 }
 
@@ -96,12 +73,6 @@ VectorMemIndex::BuildWithDataset(const DatasetPtr& dataset, const Config& config
 
     knowhere::TimeRecorder rc("BuildWithoutIds", 1);
     index_->BuildAll(dataset, index_config);
-    rc.RecordSection("TrainAndAdd");
-
-    if (is_in_nm_list(GetIndexType())) {
-        store_raw_data(dataset);
-        rc.RecordSection("store_raw_data");
-    }
     rc.ElapseFromBegin("Done");
     SetDim(index_->Dim());
 }
@@ -111,13 +82,6 @@ VectorMemIndex::Query(const DatasetPtr dataset, const SearchInfo& search_info, c
     //    AssertInfo(GetMetricType() == search_info.metric_type_,
     //               "Metric type of field index isn't the same with search info");
 
-    auto load_raw_data_closure = [&]() { LoadRawData(); };  // hide this pointer
-    auto index_type = GetIndexType();
-
-    if (is_in_nm_list(index_type)) {
-        std::call_once(raw_data_loaded_, load_raw_data_closure);
-    }
-
     auto num_queries = knowhere::GetDatasetRows(dataset);
     Config search_conf = search_info.search_params_;
     auto topk = search_info.topk_;
@@ -159,38 +123,6 @@ VectorMemIndex::Query(const DatasetPtr dataset, const SearchInfo& search_info, c
     return result;
 }
 
-void
-VectorMemIndex::store_raw_data(const knowhere::DatasetPtr& dataset) {
-    auto index_type = GetIndexType();
-    if (is_in_nm_list(index_type)) {
-        auto tensor = knowhere::GetDatasetTensor(dataset);
-        auto row_num = knowhere::GetDatasetRows(dataset);
-        auto dim = knowhere::GetDatasetDim(dataset);
-        int64_t data_size;
-        if (is_in_bin_list(index_type)) {
-            data_size = dim / 8 * row_num;
-        } else {
-            data_size = dim * row_num * sizeof(float);
-        }
-        raw_data_.resize(data_size);
-        memcpy(raw_data_.data(), tensor, data_size);
-    }
-}
-
-void
-VectorMemIndex::LoadRawData() {
-    auto index_type = GetIndexType();
-    if (is_in_nm_list(index_type)) {
-        auto bs = index_->Serialize(Config{});
-        auto bptr = std::make_shared<knowhere::Binary>();
-        auto deleter = [&](uint8_t*) {};  // avoid repeated deconstruction
-        bptr->data = std::shared_ptr<uint8_t[]>(static_cast<uint8_t*>(raw_data_.data()), deleter);
-        bptr->size = raw_data_.size();
-        bs.Append(RAW_DATA, bptr);
-        index_->Load(bs);
-    }
-}
-
 void
 VectorMemIndex::parse_config(Config& config) {
     auto stoi_closure = [](const std::string& s) -> int { return std::stoi(s); };

internal/core/src/index/VectorMemIndex.h

@@ -47,21 +47,13 @@ class VectorMemIndex : public VectorIndex {
     std::unique_ptr<SearchResult>
     Query(const DatasetPtr dataset, const SearchInfo& search_info, const BitsetView& bitset) override;
 
- private:
-    void
-    store_raw_data(const knowhere::DatasetPtr& dataset);
-
+ protected:
     void
     parse_config(Config& config);
 
-    void
-    LoadRawData();
-
- private:
+ protected:
     Config config_;
     knowhere::VecIndexPtr index_ = nullptr;
-    std::vector<uint8_t> raw_data_;
-    std::once_flag raw_data_loaded_;
 };
 
 using VectorMemIndexPtr = std::unique_ptr<VectorMemIndex>;

internal/core/src/index/VectorMemNMIndex.cpp

@@ -0,0 +1,96 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+
+#include "common/Slice.h"
+#include "common/BitsetView.h"
+#include "index/VectorMemNMIndex.h"
+#include "log/Log.h"
+
+#include "knowhere/index/VecIndexFactory.h"
+#include "knowhere/common/Timer.h"
+#include "knowhere/index/vector_index/ConfAdapterMgr.h"
+#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
+
+namespace milvus::index {
+
+BinarySet
+VectorMemNMIndex::Serialize(const Config& config) {
+    knowhere::Config serialize_config = config;
+    parse_config(serialize_config);
+
+    auto ret = index_->Serialize(serialize_config);
+    auto deleter = [&](uint8_t*) {};  // avoid repeated deconstruction
+    auto raw_data = std::shared_ptr<uint8_t[]>(static_cast<uint8_t*>(raw_data_.data()), deleter);
+    ret.Append(RAW_DATA, raw_data, raw_data_.size());
+    milvus::Disassemble(ret);
+
+    return ret;
+}
+
+void
+VectorMemNMIndex::BuildWithDataset(const DatasetPtr& dataset, const Config& config) {
+    VectorMemIndex::BuildWithDataset(dataset, config);
+    knowhere::TimeRecorder rc("store_raw_data", 1);
+    store_raw_data(dataset);
+    rc.ElapseFromBegin("Done");
+}
+
+void
+VectorMemNMIndex::Load(const BinarySet& binary_set, const Config& config) {
+    VectorMemIndex::Load(binary_set, config);
+    if (binary_set.Contains(RAW_DATA)) {
+        std::call_once(raw_data_loaded_, [&]() { LOG_SEGCORE_INFO_C << "NM index load raw data done!"; });
+    }
+}
+
+std::unique_ptr<SearchResult>
+VectorMemNMIndex::Query(const DatasetPtr dataset, const SearchInfo& search_info, const BitsetView& bitset) {
+    auto load_raw_data_closure = [&]() { LoadRawData(); };  // hide this pointer
+    // load -> query, raw data has been loaded
+    // build -> query, this case just for test, should load raw data before query
+    std::call_once(raw_data_loaded_, load_raw_data_closure);
+
+    return VectorMemIndex::Query(dataset, search_info, bitset);
+}
+
+void
+VectorMemNMIndex::store_raw_data(const knowhere::DatasetPtr& dataset) {
+    auto index_type = GetIndexType();
+    auto tensor = knowhere::GetDatasetTensor(dataset);
+    auto row_num = knowhere::GetDatasetRows(dataset);
+    auto dim = knowhere::GetDatasetDim(dataset);
+    int64_t data_size;
+    if (is_in_bin_list(index_type)) {
+        data_size = dim / 8 * row_num;
+    } else {
+        data_size = dim * row_num * sizeof(float);
+    }
+    raw_data_.resize(data_size);
+    memcpy(raw_data_.data(), tensor, data_size);
+}
+
+void
+VectorMemNMIndex::LoadRawData() {
+    auto bs = index_->Serialize(Config{});
+    auto bptr = std::make_shared<knowhere::Binary>();
+    auto deleter = [&](uint8_t*) {};  // avoid repeated deconstruction
+    bptr->data = std::shared_ptr<uint8_t[]>(static_cast<uint8_t*>(raw_data_.data()), deleter);
+    bptr->size = raw_data_.size();
+    bs.Append(RAW_DATA, bptr);
+    index_->Load(bs);
+}
+
+}  // namespace milvus::index

internal/core/src/index/VectorMemNMIndex.h

@@ -0,0 +1,61 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "index/Utils.h"
+#include "index/VectorMemIndex.h"
+
+namespace milvus::index {
+
+class VectorMemNMIndex : public VectorMemIndex {
+ public:
+    explicit VectorMemNMIndex(const IndexType& index_type, const MetricType& metric_type, const IndexMode& index_mode)
+        : VectorMemIndex(index_type, metric_type, index_mode) {
+        AssertInfo(is_in_nm_list(index_type), "not valid nm index type");
+    }
+
+    BinarySet
+    Serialize(const Config& config) override;
+
+    void
+    BuildWithDataset(const DatasetPtr& dataset, const Config& config = {}) override;
+
+    void
+    Load(const BinarySet& binary_set, const Config& config = {}) override;
+
+    std::unique_ptr<SearchResult>
+    Query(const DatasetPtr dataset, const SearchInfo& search_info, const BitsetView& bitset) override;
+
+ private:
+    void
+    store_raw_data(const knowhere::DatasetPtr& dataset);
+
+    void
+    LoadRawData();
+
+ private:
+    std::vector<uint8_t> raw_data_;
+    std::once_flag raw_data_loaded_;
+};
+
+using VectorMemNMIndexPtr = std::unique_ptr<VectorMemNMIndex>;
+}  // namespace milvus::index

internal/core/src/segcore/FieldIndexing.cpp

@@ -16,7 +16,7 @@
 
 #include "common/SystemProperty.h"
 #include "segcore/FieldIndexing.h"
-#include "index/VectorMemIndex.h"
+#include "index/VectorMemNMIndex.h"
 
 namespace milvus::segcore {
 
@@ -33,8 +33,8 @@ VectorFieldIndexing::BuildIndexRange(int64_t ack_beg, int64_t ack_end, const Vec
     data_.grow_to_at_least(ack_end);
     for (int chunk_id = ack_beg; chunk_id < ack_end; chunk_id++) {
         const auto& chunk = source->get_chunk(chunk_id);
-        auto indexing = std::make_unique<index::VectorMemIndex>(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT,
-                                                                knowhere::metric::L2, IndexMode::MODE_CPU);
+        auto indexing = std::make_unique<index::VectorMemNMIndex>(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT,
+                                                                  knowhere::metric::L2, IndexMode::MODE_CPU);
         auto dataset = knowhere::GenDataset(source->get_size_per_chunk(), dim, chunk.data());
         indexing->BuildWithDataset(dataset, conf);
         data_[chunk_id] = std::move(indexing);

internal/core/unittest/test_utils/DataGen.h

@@ -21,7 +21,7 @@
 #include "common/Schema.h"
 #include "index/ScalarIndexSort.h"
 #include "index/StringIndexSort.h"
-#include "index/VectorMemIndex.h"
+#include "index/VectorMemNMIndex.h"
 #include "knowhere/index/VecIndex.h"
 #include "knowhere/index/vector_index/adapter/VectorAdapter.h"
 #include "query/SearchOnIndex.h"
@@ -436,8 +436,8 @@ GenVecIndexing(int64_t N, int64_t dim, const float* vec) {
                                  {knowhere::indexparam::NLIST, "1024"},
                                  {knowhere::meta::DEVICE_ID, 0}};
     auto database = knowhere::GenDataset(N, dim, vec);
-    auto indexing = std::make_unique<index::VectorMemIndex>(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT,
-                                                            knowhere::metric::L2, IndexMode::MODE_CPU);
+    auto indexing = std::make_unique<index::VectorMemNMIndex>(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT,
+                                                              knowhere::metric::L2, IndexMode::MODE_CPU);
     indexing->BuildWithDataset(database, conf);
     return indexing;
 }

internal/util/indexcgowrapper/index_test.go

@@ -199,9 +199,10 @@ func TestCIndex_Codec(t *testing.T) {
 		assert.Equal(t, err, nil)
 		err = copyIndex.Load(blobs)
 		assert.Equal(t, err, nil)
-		copyBlobs, err := copyIndex.Serialize()
-		assert.Equal(t, err, nil)
-		assert.Equal(t, len(blobs), len(copyBlobs))
+		// IVF_FLAT_NM index don't support load and serialize
+		//copyBlobs, err := copyIndex.Serialize()
+		//assert.Equal(t, err, nil)
+		//assert.Equal(t, len(blobs), len(copyBlobs))
 		// TODO: check key, value and more
 
 		err = index.Delete()