From 4a66965df4c40f11eba15840cbed0970208b7d55 Mon Sep 17 00:00:00 2001 From: xige-16 Date: Sat, 5 Nov 2022 17:33:05 +0800 Subject: [PATCH] Delete RAW_DATA copy when load IVF_FLAT index data (#20274) Signed-off-by: xige-16 Signed-off-by: xige-16 --- internal/core/src/index/CMakeLists.txt | 1 + internal/core/src/index/IndexFactory.cpp | 4 + internal/core/src/index/VectorMemIndex.cpp | 68 -------------- internal/core/src/index/VectorMemIndex.h | 12 +-- internal/core/src/index/VectorMemNMIndex.cpp | 96 ++++++++++++++++++++ internal/core/src/index/VectorMemNMIndex.h | 61 +++++++++++++ internal/core/src/segcore/FieldIndexing.cpp | 6 +- internal/core/unittest/test_utils/DataGen.h | 6 +- internal/util/indexcgowrapper/index_test.go | 7 +- 9 files changed, 174 insertions(+), 87 deletions(-) create mode 100644 internal/core/src/index/VectorMemNMIndex.cpp create mode 100644 internal/core/src/index/VectorMemNMIndex.h diff --git a/internal/core/src/index/CMakeLists.txt b/internal/core/src/index/CMakeLists.txt index d9c131f3ab..11da14d5bd 100644 --- a/internal/core/src/index/CMakeLists.txt +++ b/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" ) diff --git a/internal/core/src/index/IndexFactory.cpp b/internal/core/src/index/IndexFactory.cpp index 1f7bbff211..193941356c 100644 --- a/internal/core/src/index/IndexFactory.cpp +++ b/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(index_type, metric_type, index_mode); + } // create mem index return std::make_unique(index_type, metric_type, index_mode); } diff --git a/internal/core/src/index/VectorMemIndex.cpp b/internal/core/src/index/VectorMemIndex.cpp index 22265247ad..9595fc84fa 100644 --- a/internal/core/src/index/VectorMemIndex.cpp +++ b/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(static_cast(raw_data_.data()), deleter); - - // std::shared_ptr raw_data(new uint8_t[raw_data_.size()], std::default_delete()); - // 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(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(); - auto deleter = [&](uint8_t*) {}; // avoid repeated deconstruction - bptr->data = std::shared_ptr(static_cast(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); }; diff --git a/internal/core/src/index/VectorMemIndex.h b/internal/core/src/index/VectorMemIndex.h index 9cfb9c1c1b..bed96e4573 100644 --- a/internal/core/src/index/VectorMemIndex.h +++ b/internal/core/src/index/VectorMemIndex.h @@ -47,21 +47,13 @@ class VectorMemIndex : public VectorIndex { std::unique_ptr 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 raw_data_; - std::once_flag raw_data_loaded_; }; using VectorMemIndexPtr = std::unique_ptr; diff --git a/internal/core/src/index/VectorMemNMIndex.cpp b/internal/core/src/index/VectorMemNMIndex.cpp new file mode 100644 index 0000000000..e8d2e19e7c --- /dev/null +++ b/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(static_cast(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 +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(); + auto deleter = [&](uint8_t*) {}; // avoid repeated deconstruction + bptr->data = std::shared_ptr(static_cast(raw_data_.data()), deleter); + bptr->size = raw_data_.size(); + bs.Append(RAW_DATA, bptr); + index_->Load(bs); +} + +} // namespace milvus::index diff --git a/internal/core/src/index/VectorMemNMIndex.h b/internal/core/src/index/VectorMemNMIndex.h new file mode 100644 index 0000000000..fd166e5cba --- /dev/null +++ b/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 +#include +#include +#include + +#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 + 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 raw_data_; + std::once_flag raw_data_loaded_; +}; + +using VectorMemNMIndexPtr = std::unique_ptr; +} // namespace milvus::index diff --git a/internal/core/src/segcore/FieldIndexing.cpp b/internal/core/src/segcore/FieldIndexing.cpp index 7c0c97e0f4..2dad773a5e 100644 --- a/internal/core/src/segcore/FieldIndexing.cpp +++ b/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(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT, - knowhere::metric::L2, IndexMode::MODE_CPU); + auto indexing = std::make_unique(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); diff --git a/internal/core/unittest/test_utils/DataGen.h b/internal/core/unittest/test_utils/DataGen.h index 5701b49de3..444860d6d6 100644 --- a/internal/core/unittest/test_utils/DataGen.h +++ b/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(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT, - knowhere::metric::L2, IndexMode::MODE_CPU); + auto indexing = std::make_unique(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT, + knowhere::metric::L2, IndexMode::MODE_CPU); indexing->BuildWithDataset(database, conf); return indexing; } diff --git a/internal/util/indexcgowrapper/index_test.go b/internal/util/indexcgowrapper/index_test.go index 15a1115ce7..632361eac6 100644 --- a/internal/util/indexcgowrapper/index_test.go +++ b/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()