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
35namespace proxima {
36namespace be {
37namespace index {
38
39class PersistSegment;
40using 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 */
48class 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