1 | /** |
2 | * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved. |
3 | * |
4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
5 | * you may not use this file except in compliance with the License. |
6 | * You may obtain a copy of the License at |
7 | * |
8 | * http://www.apache.org/licenses/LICENSE-2.0 |
9 | * |
10 | * Unless required by applicable law or agreed to in writing, software |
11 | * distributed under the License is distributed on an "AS IS" BASIS, |
12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
13 | * See the License for the specific language governing permissions and |
14 | * limitations under the License. |
15 | |
16 | * \author Hechong.xyf |
17 | * \date Oct 2019 |
18 | * \brief Interface of AiTheta Index Document |
19 | */ |
20 | |
21 | #ifndef __AITHETA2_INDEX_DOCUMENT_H__ |
22 | #define __AITHETA2_INDEX_DOCUMENT_H__ |
23 | |
24 | #include <cstdint> |
25 | #include <ailego/container/heap.h> |
26 | |
27 | namespace aitheta2 { |
28 | |
29 | /*! Index Document |
30 | */ |
31 | class IndexDocument { |
32 | public: |
33 | //! Constructor |
34 | IndexDocument() = default; |
35 | |
36 | //! Constructor |
37 | IndexDocument(uint64_t k, float v) : key_(k), score_(v) {} |
38 | |
39 | //! Constructor |
40 | IndexDocument(uint64_t k, float v, uint32_t i) |
41 | : key_(k), score_(v), index_(i) {} |
42 | |
43 | //! Constructor |
44 | IndexDocument(const IndexDocument &rhs) |
45 | : key_(rhs.key_), score_(rhs.score_), index_(rhs.index_) {} |
46 | |
47 | //! Assignment |
48 | IndexDocument &operator=(const IndexDocument &rhs) { |
49 | key_ = rhs.key_; |
50 | score_ = rhs.score_; |
51 | index_ = rhs.index_; |
52 | return *this; |
53 | } |
54 | |
55 | //! Less than |
56 | bool operator<(const IndexDocument &rhs) const { |
57 | return (this->score_ < rhs.score_); |
58 | } |
59 | |
60 | //! Greater than |
61 | bool operator>(const IndexDocument &rhs) const { |
62 | return (this->score_ > rhs.score_); |
63 | } |
64 | |
65 | //! Retrieve primary key |
66 | uint64_t key(void) const { |
67 | return key_; |
68 | } |
69 | |
70 | //! Retrieve score value |
71 | float score(void) const { |
72 | return score_; |
73 | } |
74 | |
75 | //! Retrieve index id |
76 | uint32_t index(void) const { |
77 | return index_; |
78 | } |
79 | |
80 | //! Retrieve mutable primary key |
81 | uint64_t *mutable_key(void) { |
82 | return &key_; |
83 | } |
84 | |
85 | //! Retrieve mutable score value |
86 | float *mutable_score(void) { |
87 | return &score_; |
88 | } |
89 | |
90 | //! Retrieve mutable index id |
91 | uint32_t *mutable_index(void) { |
92 | return &index_; |
93 | } |
94 | |
95 | //! Retrieve primary key |
96 | void set_key(uint64_t val) { |
97 | key_ = val; |
98 | } |
99 | |
100 | //! Retrieve score value |
101 | void set_score(float val) { |
102 | score_ = val; |
103 | } |
104 | |
105 | //! Retrieve index id |
106 | void set_index(uint32_t val) { |
107 | index_ = val; |
108 | } |
109 | |
110 | private: |
111 | //! Data members |
112 | uint64_t key_{0u}; |
113 | float score_{0.0f}; |
114 | uint32_t index_{0u}; |
115 | }; |
116 | |
117 | /*! Index Document Heap |
118 | */ |
119 | class IndexDocumentHeap : public ailego::Heap<IndexDocument> { |
120 | public: |
121 | //! Constructor |
122 | IndexDocumentHeap(void) : ailego::Heap<IndexDocument>() {} |
123 | |
124 | //! Constructor |
125 | IndexDocumentHeap(size_t max) : ailego::Heap<IndexDocument>(max) {} |
126 | |
127 | //! Constructor |
128 | IndexDocumentHeap(size_t max, float val) |
129 | : ailego::Heap<IndexDocument>(max), threshold_(val) {} |
130 | |
131 | //! Constructor |
132 | IndexDocumentHeap(const IndexDocumentHeap &rhs) |
133 | : ailego::Heap<IndexDocument>(rhs), threshold_(rhs.threshold_) {} |
134 | |
135 | //! Constructor |
136 | IndexDocumentHeap(IndexDocumentHeap &&rhs) |
137 | : ailego::Heap<IndexDocument>(std::move(rhs)), |
138 | threshold_(rhs.threshold_) {} |
139 | |
140 | //! Constructor |
141 | IndexDocumentHeap(const std::vector<IndexDocument> &rhs) |
142 | : ailego::Heap<IndexDocument>(rhs) {} |
143 | |
144 | //! Constructor |
145 | IndexDocumentHeap(std::vector<IndexDocument> &&rhs) |
146 | : ailego::Heap<IndexDocument>(std::move(rhs)) {} |
147 | |
148 | //! Insert a document into the heap |
149 | void emplace(uint64_t key, float score) { |
150 | if (score <= threshold_) { |
151 | ailego::Heap<IndexDocument>::emplace(key, score); |
152 | } |
153 | } |
154 | |
155 | //! Insert a document into the heap |
156 | void emplace(uint64_t key, float score, uint32_t index) { |
157 | if (score <= threshold_) { |
158 | ailego::Heap<IndexDocument>::emplace(key, score, index); |
159 | } |
160 | } |
161 | |
162 | //! Set threshold for RNN |
163 | void set_threshold(float val) { |
164 | threshold_ = val; |
165 | } |
166 | |
167 | //! Retrieve value of threshold for RNN |
168 | float threshold(void) const { |
169 | return threshold_; |
170 | } |
171 | |
172 | private: |
173 | //! members |
174 | float threshold_{std::numeric_limits<float>::max()}; |
175 | }; |
176 | |
177 | /*! Index Document List |
178 | */ |
179 | using IndexDocumentList = std::vector<IndexDocument>; |
180 | |
181 | } // namespace aitheta2 |
182 | |
183 | #endif // __AITHETA2_INDEX_DOCUMENT_H__ |
184 | |