From 207b854c3f331555e92d95aacba697d3ad823bd2 Mon Sep 17 00:00:00 2001 From: Cai Yudong Date: Thu, 26 Mar 2020 21:58:48 +0800 Subject: [PATCH] Caiyd 1655 gpu ivfflat delete (#1767) * update sdk Signed-off-by: yudong.cai * update segment interfaces Signed-off-by: yudong.cai * update some interfaces Signed-off-by: yudong.cai --- core/src/db/engine/ExecutionEngineImpl.cpp | 4 +-- core/src/db/insert/MemTable.cpp | 3 +- .../knowhere/index/vector_index/IndexHNSW.cpp | 3 +- .../knowhere/index/vector_index/VecIndex.h | 6 ++-- .../index/vector_index/helpers/Cloner.cpp | 33 ++++++++----------- .../src/index/thirdparty/faiss/gpu/GpuIndex.h | 3 +- .../thirdparty/faiss/gpu/GpuIndexFlat.cu | 5 ++- .../index/thirdparty/faiss/gpu/GpuIndexFlat.h | 1 + .../thirdparty/faiss/gpu/GpuIndexIVFFlat.cu | 12 +++++-- .../thirdparty/faiss/gpu/GpuIndexIVFFlat.h | 1 + .../thirdparty/faiss/gpu/GpuIndexIVFPQ.h | 1 + .../faiss/utils/ConcurrentBitset.cpp | 2 +- .../thirdparty/faiss/utils/ConcurrentBitset.h | 2 +- core/src/segment/SegmentWriter.cpp | 5 +-- core/src/segment/Vectors.cpp | 9 +++-- core/src/segment/Vectors.h | 5 ++- sdk/examples/simple/src/ClientTest.cpp | 18 +++++----- sdk/examples/simple/src/ClientTest.h | 3 ++ 18 files changed, 66 insertions(+), 50 deletions(-) diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index dbbee791b6..3864fe514c 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -644,13 +644,13 @@ ExecutionEngineImpl::BuildIndex(const std::string& location, EngineType engine_t knowhere::GenDatasetWithIds(Count(), Dimension(), from_index->GetRawVectors(), from_index->GetRawIds()); to_index->BuildAll(dataset, conf); uids = from_index->GetUids(); - from_index->GetBlacklist(blacklist); + blacklist = from_index->GetBlacklist(); } else if (bin_from_index) { auto dataset = knowhere::GenDatasetWithIds(Count(), Dimension(), bin_from_index->GetRawVectors(), bin_from_index->GetRawIds()); to_index->BuildAll(dataset, conf); uids = bin_from_index->GetUids(); - bin_from_index->GetBlacklist(blacklist); + blacklist = bin_from_index->GetBlacklist(); } #ifdef MILVUS_GPU_VERSION diff --git a/core/src/db/insert/MemTable.cpp b/core/src/db/insert/MemTable.cpp index 6080ee321d..fce8d62e5b 100644 --- a/core/src/db/insert/MemTable.cpp +++ b/core/src/db/insert/MemTable.cpp @@ -247,9 +247,8 @@ MemTable::ApplyDeletes() { for (auto& file : segment_files) { auto data_obj_ptr = cache::CpuCacheMgr::GetInstance()->GetIndex(file.location_); auto index = std::static_pointer_cast(data_obj_ptr); - faiss::ConcurrentBitsetPtr blacklist = nullptr; if (index != nullptr) { - index->GetBlacklist(blacklist); + faiss::ConcurrentBitsetPtr blacklist = index->GetBlacklist(); if (blacklist != nullptr) { indexes.emplace_back(index); blacklists.emplace_back(blacklist); diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp index 1cb94f4e40..6a1a445288 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp @@ -143,8 +143,7 @@ IndexHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config) { using P = std::pair; auto compare = [](const P& v1, const P& v2) { return v1.first < v2.first; }; - faiss::ConcurrentBitsetPtr blacklist = nullptr; - GetBlacklist(blacklist); + faiss::ConcurrentBitsetPtr blacklist = GetBlacklist(); #pragma omp parallel for for (unsigned int i = 0; i < rows; ++i) { std::vector

ret; diff --git a/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h index 71d5673570..3a0d064f03 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h @@ -78,9 +78,9 @@ class VecIndex : public Index { return nullptr; } - void - GetBlacklist(faiss::ConcurrentBitsetPtr& bitset_ptr) { - bitset_ptr = bitset_; + faiss::ConcurrentBitsetPtr + GetBlacklist() { + return bitset_; } void diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp index 62b2e121aa..2ba189c513 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp @@ -24,13 +24,21 @@ namespace milvus { namespace knowhere { namespace cloner { +void +CopyIndexData(const VecIndexPtr& dst_index, const VecIndexPtr& src_index) { + /* do real copy */ + auto uids = src_index->GetUids(); + dst_index->SetUids(uids); + + dst_index->SetBlacklist(src_index->GetBlacklist()); + dst_index->SetIndexSize(src_index->IndexSize()); +} + VecIndexPtr CopyGpuToCpu(const VecIndexPtr& index, const Config& config) { if (auto device_index = std::dynamic_pointer_cast(index)) { VecIndexPtr result = device_index->CopyGpuToCpu(config); - auto uids = index->GetUids(); - result->SetUids(uids); - result->SetIndexSize(index->IndexSize()); + CopyIndexData(result, index); return result; } else { KNOWHERE_THROW_MSG("index type is not gpuindex"); @@ -40,23 +48,11 @@ CopyGpuToCpu(const VecIndexPtr& index, const Config& config) { VecIndexPtr CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& config) { VecIndexPtr result; - auto uids = index->GetUids(); - int64_t index_size = index->IndexSize(); if (auto device_index = std::dynamic_pointer_cast(index)) { result = device_index->CopyCpuToGpu(device_id, config); - result->SetUids(uids); - result->SetIndexSize(index_size); - return result; - } - - if (auto device_index = std::dynamic_pointer_cast(index)) { + } else if (auto device_index = std::dynamic_pointer_cast(index)) { result = device_index->CopyGpuToGpu(device_id, config); - result->SetUids(uids); - result->SetIndexSize(index_size); - return result; - } - - if (auto cpu_index = std::dynamic_pointer_cast(index)) { + } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { result = cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { result = cpu_index->CopyCpuToGpu(device_id, config); @@ -68,8 +64,7 @@ CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& co KNOWHERE_THROW_MSG("this index type not support transfer to gpu"); } - result->SetUids(uids); - result->SetIndexSize(index_size); + CopyIndexData(result, index); return result; } diff --git a/core/src/index/thirdparty/faiss/gpu/GpuIndex.h b/core/src/index/thirdparty/faiss/gpu/GpuIndex.h index 44efd660d3..294c1fb703 100644 --- a/core/src/index/thirdparty/faiss/gpu/GpuIndex.h +++ b/core/src/index/thirdparty/faiss/gpu/GpuIndex.h @@ -10,6 +10,7 @@ #include #include +#include namespace faiss { namespace gpu { @@ -124,7 +125,7 @@ private: int k, float* outDistancesData, Index::idx_t* outIndicesData, - ConcurrentBitsetPtr bitset = nullptr) const; + ConcurrentBitsetPtr bitset = nullptr) const; /// Calls searchImpl_ for a single page of GPU-resident data, /// handling paging of the data and copies from the CPU diff --git a/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.cu b/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.cu index 6084933b5c..75d96cfa3b 100644 --- a/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.cu +++ b/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.cu @@ -221,10 +221,9 @@ GpuIndexFlat::searchImpl_(int n, auto bitsetDevice = toDevice(resources_, device_, nullptr, stream, {0}); data_->query(queries, bitsetDevice, k, outDistances, outIntLabels, true); } else { - auto bitsetData = bitset->bitset(); auto bitsetDevice = toDevice(resources_, device_, - const_cast(bitsetData), stream, - {(int) bitset->size()}); + const_cast(bitset->data()), stream, + {(int) bitset->size()}); data_->query(queries, bitsetDevice, k, outDistances, outIntLabels, true); } diff --git a/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.h b/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.h index 03cfdf7191..3050caa4c8 100644 --- a/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.h +++ b/core/src/index/thirdparty/faiss/gpu/GpuIndexFlat.h @@ -10,6 +10,7 @@ #include #include +#include namespace faiss { diff --git a/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.cu b/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.cu index 3030ea53af..807e192785 100644 --- a/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.cu +++ b/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.cu @@ -245,9 +245,15 @@ GpuIndexIVFFlat::searchImpl_(int n, static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch"); Tensor outLabels(const_cast(labels), {n, k}); - auto bitsetDevice = toDevice(resources_, device_, nullptr, stream, {0}); - - index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels); + if (!bitset) { + auto bitsetDevice = toDevice(resources_, device_, nullptr, stream, {0}); + index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels); + } else { + auto bitsetDevice = toDevice(resources_, device_, + const_cast(bitset->data()), stream, + {(int) bitset->size()}); + index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels); + } } diff --git a/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.h b/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.h index 876371eda4..a7328c31e3 100644 --- a/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.h +++ b/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFFlat.h @@ -9,6 +9,7 @@ #pragma once #include +#include namespace faiss { struct IndexIVFFlat; } diff --git a/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFPQ.h b/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFPQ.h index 086c9acac4..54b65980e8 100644 --- a/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFPQ.h +++ b/core/src/index/thirdparty/faiss/gpu/GpuIndexIVFPQ.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include namespace faiss { struct IndexIVFPQ; } diff --git a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp index 25f3d6a66c..f0cec5acce 100644 --- a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp +++ b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.cpp @@ -48,7 +48,7 @@ ConcurrentBitset::size() { } const uint8_t* -ConcurrentBitset::bitset() { +ConcurrentBitset::data() { return reinterpret_cast(bitset_.data()); } diff --git a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h index a447041834..4ccdc7c893 100644 --- a/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h +++ b/core/src/index/thirdparty/faiss/utils/ConcurrentBitset.h @@ -49,7 +49,7 @@ class ConcurrentBitset { size(); const uint8_t* - bitset(); + data(); private: size_t capacity_; diff --git a/core/src/segment/SegmentWriter.cpp b/core/src/segment/SegmentWriter.cpp index 2136e229c0..8c2766075e 100644 --- a/core/src/segment/SegmentWriter.cpp +++ b/core/src/segment/SegmentWriter.cpp @@ -247,13 +247,14 @@ SegmentWriter::Merge(const std::string& dir_to_merge, const std::string& name) { size_t SegmentWriter::Size() { // TODO(zhiru): switch to actual directory size - size_t ret = segment_ptr_->vectors_ptr_->Size(); + size_t vectors_size = segment_ptr_->vectors_ptr_->VectorsSize(); + size_t uids_size = segment_ptr_->vectors_ptr_->UidsSize(); /* if (segment_ptr_->id_bloom_filter_ptr_) { ret += segment_ptr_->id_bloom_filter_ptr_->Size(); } */ - return ret; + return (vectors_size * sizeof(uint8_t) + uids_size * sizeof(doc_id_t)); } size_t diff --git a/core/src/segment/Vectors.cpp b/core/src/segment/Vectors.cpp index 5e520568f9..04c68d5186 100644 --- a/core/src/segment/Vectors.cpp +++ b/core/src/segment/Vectors.cpp @@ -141,8 +141,13 @@ Vectors::GetCodeLength() const { } size_t -Vectors::Size() { - return data_.size() + uids_.size() * sizeof(doc_id_t); +Vectors::VectorsSize() { + return data_.size(); +} + +size_t +Vectors::UidsSize() { + return uids_.size(); } void diff --git a/core/src/segment/Vectors.h b/core/src/segment/Vectors.h index f3ad30a11f..2be6e62646 100644 --- a/core/src/segment/Vectors.h +++ b/core/src/segment/Vectors.h @@ -63,7 +63,10 @@ class Vectors { Erase(std::vector& offsets); size_t - Size(); + VectorsSize(); + + size_t + UidsSize(); void Clear(); diff --git a/sdk/examples/simple/src/ClientTest.cpp b/sdk/examples/simple/src/ClientTest.cpp index fb2e2b138d..0f5f0ffbdd 100644 --- a/sdk/examples/simple/src/ClientTest.cpp +++ b/sdk/examples/simple/src/ClientTest.cpp @@ -30,7 +30,7 @@ constexpr int64_t BATCH_ENTITY_COUNT = 100000; constexpr int64_t NQ = 5; constexpr int64_t TOP_K = 10; constexpr int64_t NPROBE = 32; -constexpr int64_t SEARCH_TARGET = 5000; // change this value, result is different +constexpr int64_t SEARCH_TARGET = BATCH_ENTITY_COUNT / 2; // change this value, result is different constexpr int64_t ADD_ENTITY_LOOP = 5; constexpr milvus::IndexType INDEX_TYPE = milvus::IndexType::IVFSQ8; constexpr int32_t NLIST = 16384; @@ -180,10 +180,18 @@ ClientTest::CreateIndex(const std::string& collection_name, milvus::IndexType ty void ClientTest::PreloadCollection(const std::string& collection_name) { + milvus_sdk::TimeRecorder rc("Preload"); milvus::Status stat = conn_->PreloadCollection(collection_name); std::cout << "PreloadCollection function call status: " << stat.message() << std::endl; } +void +ClientTest::CompactCollection(const std::string& collection_name) { + milvus_sdk::TimeRecorder rc("Compact"); + milvus::Status stat = conn_->CompactCollection(collection_name); + std::cout << "CompactCollection function call status: " << stat.message() << std::endl; +} + void ClientTest::DeleteByIds(const std::string& collection_name, const std::vector& id_array) { std::cout << "Delete entity: "; @@ -200,13 +208,6 @@ ClientTest::DeleteByIds(const std::string& collection_name, const std::vectorFlushCollection(collection_name); std::cout << "FlushCollection function call status: " << stat.message() << std::endl; } - - { - // compact table - milvus_sdk::TimeRecorder rc1("Compact"); - stat = conn_->CompactCollection(collection_name); - std::cout << "CompactCollection function call status: " << stat.message() << std::endl; - } } void @@ -255,6 +256,7 @@ ClientTest::Test() { std::vector delete_ids = {search_id_array_[0], search_id_array_[1]}; DeleteByIds(collection_name, delete_ids); + CompactCollection(collection_name); SearchEntities(collection_name, TOP_K, NPROBE); // this line get two search error since we delete two entities DropIndex(collection_name); diff --git a/sdk/examples/simple/src/ClientTest.h b/sdk/examples/simple/src/ClientTest.h index 2e676353eb..b2a2621019 100644 --- a/sdk/examples/simple/src/ClientTest.h +++ b/sdk/examples/simple/src/ClientTest.h @@ -66,6 +66,9 @@ class ClientTest { void PreloadCollection(const std::string&); + void + CompactCollection(const std::string&); + void DeleteByIds(const std::string&, const std::vector& id_array);