MemManager.h 2.0 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;
}

X
Xu Peng 已提交
22
class ExecutionEngine;
23

24 25
class MemVectors {
public:
X
Xu Peng 已提交
26 27
    explicit MemVectors(const std::shared_ptr<meta::Meta>&,
            const meta::GroupFileSchema&, const Options&);
28

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

    size_t total() const;

    size_t approximate_size() const;

35
    Status serialize(std::string& group_id);
36 37 38

    ~MemVectors();

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

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

X
Xu Peng 已提交
46 47 48
    std::shared_ptr<meta::Meta> pMeta_;
    Options options_;
    meta::GroupFileSchema schema_;
49
    IDGenerator* _pIdGenerator;
X
Xu Peng 已提交
50
    std::shared_ptr<ExecutionEngine> pEE_;
51 52 53 54

}; // MemVectors


X
Xu Peng 已提交
55
typedef std::shared_ptr<MemVectors> VectorsPtr;
56

57 58
class MemManager {
public:
X
Xu Peng 已提交
59 60
    MemManager(const std::shared_ptr<meta::Meta>& meta_, const Options& options)
        : _pMeta(meta_), options_(options) {}
61

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

64 65 66
    Status add_vectors(const std::string& group_id_,
            size_t n_, const float* vectors_, IDNumbers& vector_ids_);

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

69
private:
70 71
    Status add_vectors_no_lock(const std::string& group_id_,
            size_t n_, const float* vectors_, IDNumbers& vector_ids_);
X
Xu Peng 已提交
72
    Status mark_memory_as_immutable();
73

X
Xu Peng 已提交
74 75
    typedef std::map<std::string, VectorsPtr> MemMap;
    typedef std::vector<VectorsPtr> ImmMemPool;
76
    MemMap _memMap;
X
Xu Peng 已提交
77
    ImmMemPool _immMems;
78
    std::shared_ptr<meta::Meta> _pMeta;
X
Xu Peng 已提交
79
    Options options_;
X
Xu Peng 已提交
80
    std::mutex _mutex;
X
Xu Peng 已提交
81
    std::mutex serialization_mtx_;
82 83 84
}; // MemManager


X
Xu Peng 已提交
85 86 87
} // namespace engine
} // namespace vecwise
} // namespace zilliz
88 89

#endif