Fix pk2offset occupy too much memory (#18325)

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

Signed-off-by: xige-16 <xi.ge@zilliz.com>
pull/18541/head
xige-16 2022-08-11 10:48:38 +08:00 committed by GitHub
parent 0e95f8a690
commit dbb96a8b70
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 2 deletions

View File

@ -21,6 +21,7 @@
#include <string>
#include <utility>
#include <vector>
#include <unordered_map>
#include <tbb/concurrent_unordered_map.h>
#include <tbb/concurrent_unordered_set.h>
#include <boost/align/aligned_allocator.hpp>
@ -73,7 +74,7 @@ using InsertData = proto::segcore::InsertRecord;
using PkType = std::variant<std::monostate, int64_t, std::string>;
// tbb::concurrent_unordered_multimap equal_range too slow when multi repeated key
// using Pk2OffsetType = tbb::concurrent_unordered_multimap<PkType, int64_t, std::hash<PkType>>;
using Pk2OffsetType = tbb::concurrent_unordered_map<PkType, tbb::concurrent_unordered_set<int64_t>, std::hash<PkType>>;
using Pk2OffsetType = std::unordered_map<PkType, std::vector<int64_t>, std::hash<PkType>>;
inline bool
IsPrimaryKeyDataType(DataType data_type) {

View File

@ -42,6 +42,7 @@ struct InsertRecord {
std::vector<SegOffset>
search_pk(const PkType pk, Timestamp timestamp) const {
std::shared_lock lck(shared_mutex_);
std::vector<SegOffset> res_offsets;
auto offset_iter = pk2offset_.find(pk);
if (offset_iter != pk2offset_.end()) {
@ -57,6 +58,7 @@ struct InsertRecord {
std::vector<SegOffset>
search_pk(const PkType pk, int64_t insert_barrier) const {
std::shared_lock lck(shared_mutex_);
std::vector<SegOffset> res_offsets;
auto offset_iter = pk2offset_.find(pk);
if (offset_iter != pk2offset_.end()) {
@ -72,11 +74,13 @@ struct InsertRecord {
void
insert_pk(const PkType pk, int64_t offset) {
pk2offset_[pk].insert(offset);
std::lock_guard lck(shared_mutex_);
pk2offset_[pk].emplace_back(offset);
}
bool
empty_pks() const {
std::shared_lock lck(shared_mutex_);
return pk2offset_.empty();
}
@ -133,6 +137,7 @@ struct InsertRecord {
private:
// std::vector<std::unique_ptr<VectorBase>> fields_data_;
std::unordered_map<FieldId, std::unique_ptr<VectorBase>> fields_data_;
mutable std::shared_mutex shared_mutex_;
};
} // namespace milvus::segcore