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 Searcher |
19 | */ |
20 | |
21 | #ifndef __AITHETA2_INDEX_SEARCHER_H__ |
22 | #define __AITHETA2_INDEX_SEARCHER_H__ |
23 | |
24 | #include "index_container.h" |
25 | #include "index_context.h" |
26 | #include "index_measure.h" |
27 | #include "index_provider.h" |
28 | #include "index_stats.h" |
29 | |
30 | namespace aitheta2 { |
31 | |
32 | /*! Index Searcher |
33 | */ |
34 | class IndexSearcher : public IndexModule { |
35 | public: |
36 | //! Index Searcher Pointer |
37 | typedef std::shared_ptr<IndexSearcher> Pointer; |
38 | |
39 | /*! Index Searcher Stats |
40 | */ |
41 | class Stats : public IndexStats { |
42 | public: |
43 | //! Set count of documents loaded |
44 | void set_loaded_count(size_t count) { |
45 | loaded_count_ = count; |
46 | } |
47 | |
48 | //! Set time cost of documents loaded |
49 | void set_loaded_costtime(uint64_t cost) { |
50 | loaded_costtime_ = cost; |
51 | } |
52 | |
53 | //! Retrieve count of documents loaded |
54 | size_t loaded_count(void) const { |
55 | return loaded_count_; |
56 | } |
57 | |
58 | //! Retrieve time cost of documents loaded |
59 | uint64_t loaded_costtime(void) const { |
60 | return loaded_costtime_; |
61 | } |
62 | |
63 | //! Retrieve count of documents loaded (mutable) |
64 | size_t *mutable_loaded_count(void) { |
65 | return &loaded_count_; |
66 | } |
67 | |
68 | //! Retrieve time cost of documents loaded (mutable) |
69 | uint64_t *mutable_loaded_costtime(void) { |
70 | return &loaded_costtime_; |
71 | } |
72 | |
73 | private: |
74 | //! Members |
75 | size_t loaded_count_{0u}; |
76 | uint64_t loaded_costtime_{0u}; |
77 | }; |
78 | |
79 | /*! Index Searcher Context |
80 | */ |
81 | struct Context : public IndexContext {}; |
82 | |
83 | /*! Index Searcher Provider |
84 | */ |
85 | struct Provider : public IndexProvider {}; |
86 | |
87 | //! Destructor |
88 | virtual ~IndexSearcher(void) {} |
89 | |
90 | //! Initialize Searcher |
91 | virtual int init(const IndexParams ¶ms) = 0; |
92 | |
93 | //! Cleanup Searcher |
94 | virtual int cleanup(void) = 0; |
95 | |
96 | //! Load index from container |
97 | virtual int load(IndexContainer::Pointer cntr, |
98 | IndexMeasure::Pointer measure) = 0; |
99 | |
100 | //! Unload index |
101 | virtual int unload(void) = 0; |
102 | |
103 | //! Similarity brute force search |
104 | virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, |
105 | Context::Pointer &context) const = 0; |
106 | |
107 | //! Similarity search |
108 | virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, |
109 | Context::Pointer &context) const = 0; |
110 | |
111 | //! Similarity brute force search |
112 | virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, |
113 | uint32_t count, |
114 | Context::Pointer &context) const = 0; |
115 | |
116 | //! Similarity search |
117 | virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, |
118 | uint32_t count, Context::Pointer &context) const = 0; |
119 | |
120 | //! Retrieve statistics |
121 | virtual const Stats &stats(void) const = 0; |
122 | |
123 | //! Retrieve meta of index |
124 | virtual const IndexMeta &meta(void) const = 0; |
125 | |
126 | //! Retrieve params of index |
127 | virtual const IndexParams ¶ms(void) const = 0; |
128 | |
129 | //! Create a searcher context |
130 | virtual Context::Pointer create_context(void) const = 0; |
131 | |
132 | //! Create a searcher provider |
133 | virtual Provider::Pointer create_provider(void) const { |
134 | return Provider::Pointer(); |
135 | } |
136 | |
137 | //! Similarity search (FP16) |
138 | template <IndexMeta::FeatureTypes FT, |
139 | typename = typename std::enable_if<FT == IndexMeta::FT_FP16>::type> |
140 | int search_bf(const ailego::Float16 *vec, size_t dim, |
141 | Context::Pointer &context) const { |
142 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), context); |
143 | } |
144 | |
145 | //! Similarity search (FP32) |
146 | template <IndexMeta::FeatureTypes FT, |
147 | typename = typename std::enable_if<FT == IndexMeta::FT_FP32>::type> |
148 | int search_bf(const float *vec, size_t dim, Context::Pointer &context) const { |
149 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), context); |
150 | } |
151 | |
152 | //! Similarity search (INT8) |
153 | template <IndexMeta::FeatureTypes FT, |
154 | typename = typename std::enable_if<FT == IndexMeta::FT_INT8>::type> |
155 | int search_bf(const int8_t *vec, size_t dim, |
156 | Context::Pointer &context) const { |
157 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), context); |
158 | } |
159 | |
160 | //! Similarity search (INT4) |
161 | template <IndexMeta::FeatureTypes FT, |
162 | typename = typename std::enable_if<FT == IndexMeta::FT_INT4>::type> |
163 | int search_bf(const uint8_t *vec, size_t dim, |
164 | Context::Pointer &context) const { |
165 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), context); |
166 | } |
167 | |
168 | //! Similarity search (BINARY) |
169 | template <IndexMeta::FeatureTypes FT, typename = typename std::enable_if< |
170 | FT == IndexMeta::FT_BINARY32>::type> |
171 | int search_bf(const uint32_t *vec, size_t dim, |
172 | Context::Pointer &context) const { |
173 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), context); |
174 | } |
175 | |
176 | //! Similarity search in batch (FP16) |
177 | template <IndexMeta::FeatureTypes FT, |
178 | typename = typename std::enable_if<FT == IndexMeta::FT_FP16>::type> |
179 | int search_bf(const ailego::Float16 *vec, size_t dim, size_t rows, |
180 | Context::Pointer &context) const { |
181 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
182 | } |
183 | |
184 | //! Similarity search in batch (FP32) |
185 | template <IndexMeta::FeatureTypes FT, |
186 | typename = typename std::enable_if<FT == IndexMeta::FT_FP32>::type> |
187 | int search_bf(const float *vec, size_t dim, size_t rows, |
188 | Context::Pointer &context) const { |
189 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
190 | } |
191 | |
192 | //! Similarity search in batch (INT8) |
193 | template <IndexMeta::FeatureTypes FT, |
194 | typename = typename std::enable_if<FT == IndexMeta::FT_INT8>::type> |
195 | int search_bf(const int8_t *vec, size_t dim, size_t rows, |
196 | Context::Pointer &context) const { |
197 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
198 | } |
199 | |
200 | //! Similarity search in batch (INT4) |
201 | template <IndexMeta::FeatureTypes FT, |
202 | typename = typename std::enable_if<FT == IndexMeta::FT_INT4>::type> |
203 | int search_bf(const uint8_t *vec, size_t dim, size_t rows, |
204 | Context::Pointer &context) const { |
205 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
206 | } |
207 | |
208 | //! Similarity Search in batch (BINARY) |
209 | template <IndexMeta::FeatureTypes FT, typename = typename std::enable_if< |
210 | FT == IndexMeta::FT_BINARY32>::type> |
211 | int search_bf(const uint32_t *vec, size_t dim, size_t rows, |
212 | Context::Pointer &context) const { |
213 | return this->search_bf_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
214 | } |
215 | |
216 | //! Similarity search (FP16) |
217 | template <IndexMeta::FeatureTypes FT, |
218 | typename = typename std::enable_if<FT == IndexMeta::FT_FP16>::type> |
219 | int search(const ailego::Float16 *vec, size_t dim, |
220 | Context::Pointer &context) const { |
221 | return this->search_impl(vec, IndexQueryMeta(FT, dim), context); |
222 | } |
223 | |
224 | //! Similarity search (FP32) |
225 | template <IndexMeta::FeatureTypes FT, |
226 | typename = typename std::enable_if<FT == IndexMeta::FT_FP32>::type> |
227 | int search(const float *vec, size_t dim, Context::Pointer &context) const { |
228 | return this->search_impl(vec, IndexQueryMeta(FT, dim), context); |
229 | } |
230 | |
231 | //! Similarity search (INT8) |
232 | template <IndexMeta::FeatureTypes FT, |
233 | typename = typename std::enable_if<FT == IndexMeta::FT_INT8>::type> |
234 | int search(const int8_t *vec, size_t dim, Context::Pointer &context) const { |
235 | return this->search_impl(vec, IndexQueryMeta(FT, dim), context); |
236 | } |
237 | |
238 | //! Similarity search (INT4) |
239 | template <IndexMeta::FeatureTypes FT, |
240 | typename = typename std::enable_if<FT == IndexMeta::FT_INT4>::type> |
241 | int search(const uint8_t *vec, size_t dim, Context::Pointer &context) const { |
242 | return this->search_impl(vec, IndexQueryMeta(FT, dim), context); |
243 | } |
244 | |
245 | //! Similarity search (BINARY32) |
246 | template <IndexMeta::FeatureTypes FT, typename = typename std::enable_if< |
247 | FT == IndexMeta::FT_BINARY32>::type> |
248 | int search(const uint32_t *vec, size_t dim, Context::Pointer &context) const { |
249 | return this->search_impl(vec, IndexQueryMeta(FT, dim), context); |
250 | } |
251 | |
252 | //! Similarity search in batch (FP16) |
253 | template <IndexMeta::FeatureTypes FT, |
254 | typename = typename std::enable_if<FT == IndexMeta::FT_FP16>::type> |
255 | int search(const ailego::Float16 *vec, size_t dim, size_t rows, |
256 | Context::Pointer &context) const { |
257 | return this->search_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
258 | } |
259 | |
260 | //! Similarity search in batch (FP32) |
261 | template <IndexMeta::FeatureTypes FT, |
262 | typename = typename std::enable_if<FT == IndexMeta::FT_FP32>::type> |
263 | int search(const float *vec, size_t dim, size_t rows, |
264 | Context::Pointer &context) const { |
265 | return this->search_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
266 | } |
267 | |
268 | //! Similarity search in batch (INT8) |
269 | template <IndexMeta::FeatureTypes FT, |
270 | typename = typename std::enable_if<FT == IndexMeta::FT_INT8>::type> |
271 | int search(const int8_t *vec, size_t dim, size_t rows, |
272 | Context::Pointer &context) const { |
273 | return this->search_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
274 | } |
275 | |
276 | //! Similarity Search in batch (INT4) |
277 | template <IndexMeta::FeatureTypes FT, |
278 | typename = typename std::enable_if<FT == IndexMeta::FT_INT4>::type> |
279 | int search(const uint8_t *vec, size_t dim, size_t rows, |
280 | Context::Pointer &context) const { |
281 | return this->search_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
282 | } |
283 | |
284 | //! Similarity Search in batch (BINARY) |
285 | template <IndexMeta::FeatureTypes FT, typename = typename std::enable_if< |
286 | FT == IndexMeta::FT_BINARY32>::type> |
287 | int search(const uint32_t *vec, size_t dim, size_t rows, |
288 | Context::Pointer &context) const { |
289 | return this->search_impl(vec, IndexQueryMeta(FT, dim), rows, context); |
290 | } |
291 | }; |
292 | |
293 | } // namespace aitheta2 |
294 | |
295 | #endif // __AITHETA2_INDEX_SEARCHER_H__ |
296 | |