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 Hongqing.hu |
17 | * \date Feb 2021 |
18 | * \brief Write request interface definition for proxima search engine |
19 | */ |
20 | |
21 | #pragma once |
22 | |
23 | #include "index/collection_dataset.h" |
24 | |
25 | namespace proxima { |
26 | namespace be { |
27 | namespace agent { |
28 | |
29 | class WriteRequest; |
30 | using WriteRequestPtr = std::shared_ptr<WriteRequest>; |
31 | |
32 | /*! Write Request |
33 | */ |
34 | class WriteRequest { |
35 | public: |
36 | //! Request type |
37 | enum class RequestType : uint32_t { PROXY = 0, DIRECT = 1 }; |
38 | |
39 | //! Constructor |
40 | WriteRequest() = default; |
41 | |
42 | //! Destructor |
43 | ~WriteRequest() = default; |
44 | |
45 | public: |
46 | //! Set request type |
47 | void set_request_type(RequestType type) { |
48 | request_type_ = type; |
49 | } |
50 | |
51 | //! Get request type |
52 | RequestType request_type() const { |
53 | return request_type_; |
54 | } |
55 | |
56 | // Is proxy request |
57 | bool is_proxy_request() const { |
58 | return request_type_ == RequestType::PROXY; |
59 | } |
60 | |
61 | //! Set magic number |
62 | void set_magic_number(uint64_t number) { |
63 | magic_number_ = number; |
64 | } |
65 | |
66 | //! Get magic number |
67 | uint64_t magic_number() const { |
68 | return magic_number_; |
69 | } |
70 | |
71 | //! Set collection name |
72 | void set_collection_name(const std::string &collection) { |
73 | collection_name_ = collection; |
74 | } |
75 | |
76 | //! Get collection name |
77 | const std::string &collection_name() const { |
78 | return collection_name_; |
79 | } |
80 | |
81 | //! Add collection dataset |
82 | void add_collection_dataset(index::CollectionDatasetPtr dataset) { |
83 | row_count_ += dataset->size(); |
84 | records_.emplace_back(std::move(dataset)); |
85 | } |
86 | |
87 | //! Get collection dataset |
88 | const index::CollectionDatasetPtr &get_collection_dataset( |
89 | size_t index) const { |
90 | return records_[index]; |
91 | } |
92 | |
93 | //! Get all collection dataset |
94 | const std::vector<index::CollectionDatasetPtr> &get_collection_dataset() |
95 | const { |
96 | return records_; |
97 | } |
98 | |
99 | //! Get collection dataset count |
100 | size_t collection_dataset_count() const { |
101 | return records_.size(); |
102 | } |
103 | |
104 | //! Get rows count |
105 | size_t row_count() const { |
106 | return row_count_; |
107 | } |
108 | |
109 | private: |
110 | //! Members |
111 | RequestType request_type_{RequestType::PROXY}; |
112 | size_t row_count_{0}; |
113 | uint64_t magic_number_{0}; |
114 | std::string collection_name_{}; |
115 | std::vector<index::CollectionDatasetPtr> records_{}; |
116 | }; |
117 | |
118 | } // end namespace agent |
119 | } // namespace be |
120 | } // end namespace proxima |
121 | |