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
30namespace aitheta2 {
31
32/*! Index Searcher
33 */
34class 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 &params) = 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 &params(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