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 guonix |
17 | * \date Nov 2020 |
18 | * \brief |
19 | */ |
20 | |
21 | #pragma once |
22 | |
23 | #include <functional> // for std::reference_wrapper |
24 | #include "collection_query.h" |
25 | #include "knn_task.h" |
26 | |
27 | namespace proxima { |
28 | namespace be { |
29 | namespace query { |
30 | |
31 | /*! |
32 | * KNNQuery: handler of knn query |
33 | */ |
34 | class KNNQuery : public CollectionQuery, public KNNQueryContext { |
35 | public: |
36 | // Alias for reference to Result |
37 | using ResultRef = std::reference_wrapper<const index::QueryResult>; |
38 | |
39 | // Alias for result reference list, which used to merge result and sort |
40 | using ResultRefList = std::vector<ResultRef>; |
41 | |
42 | // Comparator for ResultReference |
43 | struct ResultRefCompare { |
44 | bool operator()(const ResultRef &__x, const ResultRef &__y) const { |
45 | return __x.get() < __y.get(); |
46 | } |
47 | }; |
48 | |
49 | // Heap of ResultReference |
50 | using ResultRefHeap = ailego::Heap<ResultRef, ResultRefCompare>; |
51 | |
52 | public: |
53 | //! Constructor |
54 | KNNQuery(uint64_t traceID, const proto::QueryRequest *req, |
55 | index::IndexServicePtr index, MetaWrapperPtr meta_wrapper, |
56 | ExecutorPtr executor_ptr, ProfilerPtr profiler_ptr, |
57 | proto::QueryResponse *resp); |
58 | |
59 | //! Destructor |
60 | ~KNNQuery() override; |
61 | |
62 | public: |
63 | //! Validate query object, 0 for valid, otherwise non zero returned |
64 | int validate() const override; |
65 | |
66 | //! Retrieve IOMode of query |
67 | IOMode mode() const override; |
68 | |
69 | //! Retrieve the type of query, Readonly |
70 | QueryType type() const override; |
71 | |
72 | //! Prepare resources, 0 for success, otherwise failed |
73 | int prepare() override; |
74 | |
75 | //! Evaluate query, and collection feedback |
76 | int evaluate() override; |
77 | |
78 | //! Finalize query object |
79 | int finalize() override; |
80 | |
81 | //! Retrieve column name |
82 | const std::string &column() const override; |
83 | |
84 | //! Retrieve features field |
85 | const std::string &features() const override; |
86 | |
87 | //! Retrieve batch_count field |
88 | uint32_t batch_count() const override; |
89 | |
90 | //! Retrieve QueryParams |
91 | const index::QueryParams &query_params() const override; |
92 | |
93 | private: |
94 | //! Build query param from PB proto |
95 | int build_query_param(const proto::QueryRequest::KnnQueryParam &); |
96 | |
97 | //! Collect result |
98 | int collect_result(); |
99 | |
100 | //! Feed entity field |
101 | uint32_t feed_entity(const ResultRefList &, proto::QueryResponse::Result *); |
102 | |
103 | //! transform query features if necessary |
104 | int transform_feature(const proto::QueryRequest::KnnQueryParam &); |
105 | |
106 | private: |
107 | //! QueryParams handler |
108 | index::QueryParams query_param_{}; |
109 | |
110 | //! KNNTasks, which scheduled by executor |
111 | KNNTaskPtrList tasks_{}; |
112 | |
113 | //! Query features |
114 | std::string features_{}; |
115 | }; |
116 | |
117 | |
118 | } // namespace query |
119 | } // namespace be |
120 | } // namespace proxima |
121 | |