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 | #include "equal_query.h" |
22 | #include "common/error_code.h" |
23 | #include "common/logger.h" |
24 | |
25 | namespace proxima { |
26 | namespace be { |
27 | namespace query { |
28 | |
29 | //! Constructor |
30 | EqualQuery::EqualQuery(uint64_t traceID, const proto::GetDocumentRequest *req, |
31 | index::IndexServicePtr index, |
32 | MetaWrapperPtr meta_wrapper, ExecutorPtr executor_ptr, |
33 | ProfilerPtr profiler_ptr, |
34 | proto::GetDocumentResponse *resp) |
35 | : ContextImpl(traceID, index, meta_wrapper, profiler_ptr, executor_ptr), |
36 | request_(req), |
37 | response_(resp) {} |
38 | |
39 | //! Destructor |
40 | EqualQuery::~EqualQuery() = default; |
41 | |
42 | //! Validate query object, 0 for valid, otherwise non zero returned |
43 | int EqualQuery::validate() const { |
44 | // Avoid core dump resulted by nullptr request |
45 | if (!request_) { |
46 | return PROXIMA_BE_ERROR_CODE(InvalidArgument); |
47 | } |
48 | |
49 | int code = ContextImpl::validate(); |
50 | if (code == 0) { |
51 | if (valid_executor()) { |
52 | code = meta()->validate_collection(collection()); |
53 | } else { |
54 | LOG_WARN("Invalid response or executor passed to EqualQuery" ); |
55 | code = PROXIMA_BE_ERROR_CODE(InvalidArgument); |
56 | } |
57 | } |
58 | return code; |
59 | } |
60 | |
61 | //! Retrieve IOMode of query |
62 | IOMode EqualQuery::mode() const { |
63 | return IOMode::READONLY; |
64 | } |
65 | |
66 | //! Retrieve the type of query, Readonly |
67 | QueryType EqualQuery::type() const { |
68 | return QueryType::EQUAL; |
69 | } |
70 | |
71 | //! Prepare resources, 0 for success, otherwise failed |
72 | int EqualQuery::prepare() { |
73 | index::SegmentPtrList segments; |
74 | int code = list_segments(&segments); |
75 | if (code == 0) { |
76 | for (auto &segment : segments) { |
77 | tasks_.emplace_back(std::make_shared<EqualTask>(segment, this)); |
78 | } |
79 | } else { |
80 | LOG_ERROR("Failed to build query param from request" ); |
81 | } |
82 | return code; |
83 | } |
84 | |
85 | //! Evaluate query, and collection feedback |
86 | int EqualQuery::evaluate() { |
87 | TaskPtrList tasks(tasks_.begin(), tasks_.end()); |
88 | int code = executor()->execute_tasks(tasks); |
89 | if (code == 0) { |
90 | // one result or nothing |
91 | for (auto &task : tasks_) { |
92 | if (task->hit()) { // Pick first hit document, drop others |
93 | proto::Document *doc = response_->mutable_document(); |
94 | doc->set_primary_key(primary_key()); |
95 | code = fill_forward(task->forward(), doc); |
96 | if (code != 0) { |
97 | LOG_WARN("Fill forward failed. code[%d] what[%s]" , code, |
98 | ErrorCode::What(code)); |
99 | } |
100 | break; |
101 | } |
102 | } |
103 | } |
104 | return code; |
105 | } |
106 | |
107 | int EqualQuery::finalize() { |
108 | return 0; |
109 | } |
110 | |
111 | //! Retrieve collection name |
112 | const std::string &EqualQuery::collection() const { |
113 | return request_->collection_name(); |
114 | } |
115 | |
116 | //! Retrieve primary_key |
117 | uint64_t EqualQuery::primary_key() const { |
118 | return request_->primary_key(); |
119 | } |
120 | |
121 | } // namespace query |
122 | } // namespace be |
123 | } // namespace proxima |
124 | |