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 Haichao.chc |
17 | * \date Oct 2020 |
18 | * \brief Provides search ability for persist segment |
19 | */ |
20 | |
21 | #pragma once |
22 | |
23 | #include <unordered_map> |
24 | #include <ailego/parallel/lock.h> |
25 | #include "common/macro_define.h" |
26 | #include "meta/meta.h" |
27 | #include "segment.h" |
28 | #include "../column/column_reader.h" |
29 | #include "../column/forward_reader.h" |
30 | #include "../concurrent_hash_map.h" |
31 | #include "../delete_store.h" |
32 | #include "../id_map.h" |
33 | #include "../typedef.h" |
34 | |
35 | namespace proxima { |
36 | namespace be { |
37 | namespace index { |
38 | |
39 | class PersistSegment; |
40 | using PersistSegmentPtr = std::shared_ptr<PersistSegment>; |
41 | |
42 | /* |
43 | * A PersistSegment represents block of index data in persist storage. |
44 | * It transform from MemorySegment which will dump to PersistSegment |
45 | * at some condition. |
46 | * It also provides search abilify, but readonly |
47 | */ |
48 | class PersistSegment : public Segment { |
49 | public: |
50 | PROXIMA_DISALLOW_COPY_AND_ASSIGN(PersistSegment); |
51 | |
52 | //! Constructor |
53 | PersistSegment(const std::string &coll_name, const std::string &coll_path, |
54 | const SegmentMeta &seg_meta, |
55 | const meta::CollectionMeta *schema_ptr, |
56 | const DeleteStore *delete_store_ptr, const IDMap *id_map_ptr, |
57 | uint32_t concurrency_val) |
58 | : schema_(schema_ptr), |
59 | delete_store_(delete_store_ptr), |
60 | id_map_(id_map_ptr), |
61 | concurrency_(concurrency_val) { |
62 | this->set_collection_name(coll_name); |
63 | this->set_collection_path(coll_path); |
64 | this->set_segment_meta(seg_meta); |
65 | } |
66 | |
67 | //! Destructor |
68 | ~PersistSegment() override; |
69 | |
70 | //! Create an instance |
71 | static PersistSegmentPtr Create(const std::string &coll_name, |
72 | const std::string &coll_path, |
73 | const SegmentMeta &seg_meta, |
74 | const meta::CollectionMeta *schema_ptr, |
75 | const DeleteStore *delete_store_ptr, |
76 | const IDMap *id_map_ptr, |
77 | uint32_t concurrency_val); |
78 | |
79 | //! Create and load an instance |
80 | static int CreateAndLoad(const std::string &collection_name, |
81 | const std::string &collection_path, |
82 | const SegmentMeta &segment_meta, |
83 | const meta::CollectionMeta *schema, |
84 | const DeleteStore *delete_store, const IDMap *id_map, |
85 | uint32_t concurrency, |
86 | const ReadOptions &read_options, |
87 | PersistSegmentPtr *persist_segment); |
88 | |
89 | public: |
90 | //! Load index from persist storage |
91 | int load(const ReadOptions &read_options); |
92 | |
93 | //! Unload index |
94 | int unload(); |
95 | |
96 | public: |
97 | //! Knn similar search |
98 | int knn_search(const std::string &column_name, const std::string &query, |
99 | const QueryParams &query_params, |
100 | QueryResultList *results) override; |
101 | |
102 | //! Knn similar search with batch mode |
103 | int knn_search(const std::string &column_name, const std::string &query, |
104 | const QueryParams &query_params, uint32_t batch_count, |
105 | std::vector<QueryResultList> *results) override; |
106 | |
107 | //! Just search forward by doc primary key |
108 | int kv_search(uint64_t primary_key, QueryResult *result) override; |
109 | |
110 | public: |
111 | //! Remove a index column |
112 | int remove_column(const std::string &column_name) override; |
113 | |
114 | //! Add a index column |
115 | int add_column(const meta::ColumnMetaPtr &column_meta) override; |
116 | |
117 | public: |
118 | //! Return forward count |
119 | size_t doc_count() const override { |
120 | return forward_reader_->doc_count(); |
121 | } |
122 | |
123 | public: |
124 | //! Get forward reader |
125 | ForwardReaderPtr get_forward_reader() const override { |
126 | return forward_reader_; |
127 | } |
128 | |
129 | //! Get column reader |
130 | ColumnReaderPtr get_column_reader( |
131 | const std::string &column_name) const override { |
132 | if (!column_readers_.has(column_name)) { |
133 | return ColumnReaderPtr(); |
134 | } |
135 | return column_readers_.get(column_name); |
136 | } |
137 | |
138 | private: |
139 | int load_forward_reader(const ReadOptions &read_options); |
140 | |
141 | int load_column_readers(const ReadOptions &read_options); |
142 | |
143 | private: |
144 | static constexpr uint32_t MAX_WAIT_RETRY_COUNT = 60U; |
145 | |
146 | private: |
147 | const meta::CollectionMeta *schema_{nullptr}; |
148 | const DeleteStore *delete_store_{nullptr}; |
149 | const IDMap *id_map_{nullptr}; |
150 | uint32_t concurrency_{0U}; |
151 | |
152 | ForwardReaderPtr forward_reader_{}; |
153 | ConcurrentHashMap<std::string, ColumnReaderPtr> column_readers_{}; |
154 | |
155 | std::atomic<uint64_t> active_search_count_{0U}; |
156 | bool loaded_{false}; |
157 | }; |
158 | |
159 | |
160 | } // end namespace index |
161 | } // namespace be |
162 | } // end namespace proxima |
163 | |