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
25namespace proxima {
26namespace be {
27namespace query {
28
29//! Constructor
30EqualQuery::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
40EqualQuery::~EqualQuery() = default;
41
42//! Validate query object, 0 for valid, otherwise non zero returned
43int 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
62IOMode EqualQuery::mode() const {
63 return IOMode::READONLY;
64}
65
66//! Retrieve the type of query, Readonly
67QueryType EqualQuery::type() const {
68 return QueryType::EQUAL;
69}
70
71//! Prepare resources, 0 for success, otherwise failed
72int 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
86int 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
107int EqualQuery::finalize() {
108 return 0;
109}
110
111//! Retrieve collection name
112const std::string &EqualQuery::collection() const {
113 return request_->collection_name();
114}
115
116//! Retrieve primary_key
117uint64_t EqualQuery::primary_key() const {
118 return request_->primary_key();
119}
120
121} // namespace query
122} // namespace be
123} // namespace proxima
124