MemManager.h 2.6 KB
Newer Older
X
Xu Peng 已提交
1 2 3 4 5 6
/*******************************************************************************
 * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved
 * Unauthorized copying of this file, via any medium is strictly prohibited.
 * Proprietary and confidential.
 ******************************************************************************/
#pragma once
7

G
groot 已提交
8
#include "ExecutionEngine.h"
9 10 11 12
#include "IDGenerator.h"
#include "Status.h"
#include "Meta.h"

13 14
#include <map>
#include <string>
X
Xu Peng 已提交
15
#include <ctime>
X
Xu Peng 已提交
16 17
#include <memory>
#include <mutex>
S
starlord 已提交
18
#include <set>
19

X
Xu Peng 已提交
20
namespace zilliz {
J
jinhai 已提交
21
namespace milvus {
X
Xu Peng 已提交
22
namespace engine {
23

24 25 26 27
namespace meta {
    class Meta;
}

28 29
class MemVectors {
public:
X
Xu Peng 已提交
30
    using MetaPtr = meta::Meta::Ptr;
G
groot 已提交
31
    using Ptr = std::shared_ptr<MemVectors>;
X
Xu Peng 已提交
32

X
Xu Peng 已提交
33
    explicit MemVectors(const std::shared_ptr<meta::Meta>&,
34
            const meta::TableFileSchema&, const Options&);
35

S
starlord 已提交
36
    Status Add(size_t n_, const float* vectors_, IDNumbers& vector_ids_);
37

S
starlord 已提交
38
    size_t RowCount() const;
39

S
starlord 已提交
40
    size_t Size() const;
41

X
Xu Peng 已提交
42
    Status Serialize(std::string& table_id);
43 44 45

    ~MemVectors();

G
groot 已提交
46
    const std::string& Location() const { return schema_.location_; }
X
Xu Peng 已提交
47

S
starlord 已提交
48 49
    std::string TableId() const { return schema_.table_id_; }

50
private:
X
Xu Peng 已提交
51 52 53 54
    MemVectors() = delete;
    MemVectors(const MemVectors&) = delete;
    MemVectors& operator=(const MemVectors&) = delete;

S
starlord 已提交
55
    MetaPtr meta_;
X
Xu Peng 已提交
56
    Options options_;
57
    meta::TableFileSchema schema_;
S
starlord 已提交
58 59
    IDGenerator* id_generator_;
    ExecutionEnginePtr active_engine_;
60 61 62 63

}; // MemVectors


64

65 66
class MemManager {
public:
X
Xu Peng 已提交
67
    using MetaPtr = meta::Meta::Ptr;
G
groot 已提交
68 69
    using MemVectorsPtr = typename MemVectors::Ptr;
    using Ptr = std::shared_ptr<MemManager>;
70

X
Xu Peng 已提交
71
    MemManager(const std::shared_ptr<meta::Meta>& meta, const Options& options)
S
starlord 已提交
72
        : meta_(meta), options_(options) {}
73

X
Xu Peng 已提交
74
    MemVectorsPtr GetMemByTable(const std::string& table_id);
75

X
Xu Peng 已提交
76 77
    Status InsertVectors(const std::string& table_id,
            size_t n, const float* vectors, IDNumbers& vector_ids);
78

S
starlord 已提交
79
    Status Serialize(std::set<std::string>& table_ids);
X
Xu Peng 已提交
80

G
groot 已提交
81 82
    Status EraseMemVector(const std::string& table_id);

83
private:
X
Xu Peng 已提交
84 85 86 87
    Status InsertVectorsNoLock(const std::string& table_id,
            size_t n, const float* vectors, IDNumbers& vector_ids);
    Status ToImmutable();

S
starlord 已提交
88 89 90 91 92
    using MemIdMap = std::map<std::string, MemVectorsPtr>;
    using MemList = std::vector<MemVectorsPtr>;
    MemIdMap mem_id_map_;
    MemList immu_mem_list_;
    MetaPtr meta_;
X
Xu Peng 已提交
93
    Options options_;
X
Xu Peng 已提交
94
    std::mutex mutex_;
X
Xu Peng 已提交
95
    std::mutex serialization_mtx_;
96 97 98
}; // MemManager


X
Xu Peng 已提交
99
} // namespace engine
J
jinhai 已提交
100
} // namespace milvus
X
Xu Peng 已提交
101
} // namespace zilliz