MemManager.h 2.2 KB
Newer Older
1 2 3 4 5
#ifndef STORAGE_VECENGINE_MEMVECTORS_H_
#define STORAGE_VECENGINE_MEMVECTORS_H_

#include <map>
#include <string>
X
Xu Peng 已提交
6
#include <ctime>
X
Xu Peng 已提交
7 8
#include <memory>
#include <mutex>
X
Xu Peng 已提交
9 10
#include "IDGenerator.h"
#include "Status.h"
X
Xu Peng 已提交
11
#include "Meta.h"
12 13


X
Xu Peng 已提交
14 15 16
namespace zilliz {
namespace vecwise {
namespace engine {
17

18 19 20 21
namespace meta {
    class Meta;
}

22
template <typename EngineT>
23 24
class MemVectors {
public:
X
Xu Peng 已提交
25
    typedef typename EngineT::Ptr EnginePtr;
X
Xu Peng 已提交
26
    typedef typename meta::Meta::Ptr MetaPtr;
X
Xu Peng 已提交
27

X
Xu Peng 已提交
28 29
    explicit MemVectors(const std::shared_ptr<meta::Meta>&,
            const meta::GroupFileSchema&, const Options&);
30

X
Xu Peng 已提交
31
    void add(size_t n_, const float* vectors_, IDNumbers& vector_ids_);
32 33 34 35 36

    size_t total() const;

    size_t approximate_size() const;

37
    Status serialize(std::string& group_id);
38 39 40

    ~MemVectors();

X
Xu Peng 已提交
41
    const std::string& location() const { return schema_.location; }
X
Xu Peng 已提交
42

43
private:
X
Xu Peng 已提交
44 45 46 47
    MemVectors() = delete;
    MemVectors(const MemVectors&) = delete;
    MemVectors& operator=(const MemVectors&) = delete;

X
Xu Peng 已提交
48
    MetaPtr pMeta_;
X
Xu Peng 已提交
49 50
    Options options_;
    meta::GroupFileSchema schema_;
51
    IDGenerator* _pIdGenerator;
X
Xu Peng 已提交
52
    EnginePtr pEE_;
53 54 55 56

}; // MemVectors


57

58
template<typename EngineT>
59 60
class MemManager {
public:
61 62
    typedef MemVectors<EngineT> ItemT;
    typedef std::shared_ptr<ItemT> VectorsPtr;
X
Xu Peng 已提交
63
    typedef typename meta::Meta::Ptr MetaPtr;
64

X
Xu Peng 已提交
65 66
    MemManager(const std::shared_ptr<meta::Meta>& meta_, const Options& options)
        : _pMeta(meta_), options_(options) {}
67

X
Xu Peng 已提交
68
    VectorsPtr get_mem_by_group(const std::string& group_id_);
69

70 71 72
    Status add_vectors(const std::string& group_id_,
            size_t n_, const float* vectors_, IDNumbers& vector_ids_);

73
    Status serialize(std::vector<std::string>& group_ids);
X
Xu Peng 已提交
74

75
private:
76 77
    Status add_vectors_no_lock(const std::string& group_id_,
            size_t n_, const float* vectors_, IDNumbers& vector_ids_);
X
Xu Peng 已提交
78
    Status mark_memory_as_immutable();
79

X
Xu Peng 已提交
80 81
    typedef std::map<std::string, VectorsPtr> MemMap;
    typedef std::vector<VectorsPtr> ImmMemPool;
82
    MemMap _memMap;
X
Xu Peng 已提交
83
    ImmMemPool _immMems;
X
Xu Peng 已提交
84
    MetaPtr _pMeta;
X
Xu Peng 已提交
85
    Options options_;
X
Xu Peng 已提交
86
    std::mutex _mutex;
X
Xu Peng 已提交
87
    std::mutex serialization_mtx_;
88 89 90
}; // MemManager


X
Xu Peng 已提交
91 92 93
} // namespace engine
} // namespace vecwise
} // namespace zilliz
94
#include "MemManager.cpp"
95 96

#endif