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
27namespace proxima {
28namespace be {
29namespace query {
30
31/*!
32 * KNNQuery: handler of knn query
33 */
34class 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