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 Hechong.xyf
17 * \date Oct 2019
18 * \brief Interface of AiTheta Index Document
19 */
20
21#ifndef __AITHETA2_INDEX_DOCUMENT_H__
22#define __AITHETA2_INDEX_DOCUMENT_H__
23
24#include <cstdint>
25#include <ailego/container/heap.h>
26
27namespace aitheta2 {
28
29/*! Index Document
30 */
31class IndexDocument {
32 public:
33 //! Constructor
34 IndexDocument() = default;
35
36 //! Constructor
37 IndexDocument(uint64_t k, float v) : key_(k), score_(v) {}
38
39 //! Constructor
40 IndexDocument(uint64_t k, float v, uint32_t i)
41 : key_(k), score_(v), index_(i) {}
42
43 //! Constructor
44 IndexDocument(const IndexDocument &rhs)
45 : key_(rhs.key_), score_(rhs.score_), index_(rhs.index_) {}
46
47 //! Assignment
48 IndexDocument &operator=(const IndexDocument &rhs) {
49 key_ = rhs.key_;
50 score_ = rhs.score_;
51 index_ = rhs.index_;
52 return *this;
53 }
54
55 //! Less than
56 bool operator<(const IndexDocument &rhs) const {
57 return (this->score_ < rhs.score_);
58 }
59
60 //! Greater than
61 bool operator>(const IndexDocument &rhs) const {
62 return (this->score_ > rhs.score_);
63 }
64
65 //! Retrieve primary key
66 uint64_t key(void) const {
67 return key_;
68 }
69
70 //! Retrieve score value
71 float score(void) const {
72 return score_;
73 }
74
75 //! Retrieve index id
76 uint32_t index(void) const {
77 return index_;
78 }
79
80 //! Retrieve mutable primary key
81 uint64_t *mutable_key(void) {
82 return &key_;
83 }
84
85 //! Retrieve mutable score value
86 float *mutable_score(void) {
87 return &score_;
88 }
89
90 //! Retrieve mutable index id
91 uint32_t *mutable_index(void) {
92 return &index_;
93 }
94
95 //! Retrieve primary key
96 void set_key(uint64_t val) {
97 key_ = val;
98 }
99
100 //! Retrieve score value
101 void set_score(float val) {
102 score_ = val;
103 }
104
105 //! Retrieve index id
106 void set_index(uint32_t val) {
107 index_ = val;
108 }
109
110 private:
111 //! Data members
112 uint64_t key_{0u};
113 float score_{0.0f};
114 uint32_t index_{0u};
115};
116
117/*! Index Document Heap
118 */
119class IndexDocumentHeap : public ailego::Heap<IndexDocument> {
120 public:
121 //! Constructor
122 IndexDocumentHeap(void) : ailego::Heap<IndexDocument>() {}
123
124 //! Constructor
125 IndexDocumentHeap(size_t max) : ailego::Heap<IndexDocument>(max) {}
126
127 //! Constructor
128 IndexDocumentHeap(size_t max, float val)
129 : ailego::Heap<IndexDocument>(max), threshold_(val) {}
130
131 //! Constructor
132 IndexDocumentHeap(const IndexDocumentHeap &rhs)
133 : ailego::Heap<IndexDocument>(rhs), threshold_(rhs.threshold_) {}
134
135 //! Constructor
136 IndexDocumentHeap(IndexDocumentHeap &&rhs)
137 : ailego::Heap<IndexDocument>(std::move(rhs)),
138 threshold_(rhs.threshold_) {}
139
140 //! Constructor
141 IndexDocumentHeap(const std::vector<IndexDocument> &rhs)
142 : ailego::Heap<IndexDocument>(rhs) {}
143
144 //! Constructor
145 IndexDocumentHeap(std::vector<IndexDocument> &&rhs)
146 : ailego::Heap<IndexDocument>(std::move(rhs)) {}
147
148 //! Insert a document into the heap
149 void emplace(uint64_t key, float score) {
150 if (score <= threshold_) {
151 ailego::Heap<IndexDocument>::emplace(key, score);
152 }
153 }
154
155 //! Insert a document into the heap
156 void emplace(uint64_t key, float score, uint32_t index) {
157 if (score <= threshold_) {
158 ailego::Heap<IndexDocument>::emplace(key, score, index);
159 }
160 }
161
162 //! Set threshold for RNN
163 void set_threshold(float val) {
164 threshold_ = val;
165 }
166
167 //! Retrieve value of threshold for RNN
168 float threshold(void) const {
169 return threshold_;
170 }
171
172 private:
173 //! members
174 float threshold_{std::numeric_limits<float>::max()};
175};
176
177/*! Index Document List
178 */
179using IndexDocumentList = std::vector<IndexDocument>;
180
181} // namespace aitheta2
182
183#endif // __AITHETA2_INDEX_DOCUMENT_H__
184