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
25namespace proxima {
26namespace be {
27namespace agent {
28
29class WriteRequest;
30using WriteRequestPtr = std::shared_ptr<WriteRequest>;
31
32/*! Write Request
33 */
34class 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