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 | |
17 | #include "index/segment/memory_segment.h" |
18 | #include <ailego/utility/file_helper.h> |
19 | #include <gtest/gtest.h> |
20 | |
21 | using namespace proxima::be; |
22 | using namespace proxima::be::index; |
23 | |
24 | class MemorySegmentTest : public testing::Test { |
25 | protected: |
26 | void SetUp() { |
27 | ailego::FileHelper::RemoveDirectory("./teachers/" ); |
28 | FillSchema(); |
29 | } |
30 | |
31 | void TearDown() {} |
32 | |
33 | void FillSchema() { |
34 | schema_ = std::make_shared<meta::CollectionMeta>(); |
35 | meta::ColumnMetaPtr column_meta = std::make_shared<meta::ColumnMeta>(); |
36 | column_meta->set_name("face" ); |
37 | column_meta->set_index_type(IndexTypes::PROXIMA_GRAPH_INDEX); |
38 | column_meta->set_data_type(DataTypes::VECTOR_FP32); |
39 | column_meta->set_dimension(16); |
40 | column_meta->mutable_parameters()->set("metric_type" , "SquaredEuclidean" ); |
41 | schema_->append(column_meta); |
42 | schema_->set_name("teachers" ); |
43 | } |
44 | |
45 | protected: |
46 | meta::CollectionMetaPtr schema_{}; |
47 | }; |
48 | |
49 | TEST_F(MemorySegmentTest, TestGeneral) { |
50 | DeleteStore delete_store("teachers" , "./teachers/" ); |
51 | ReadOptions read_options; |
52 | read_options.use_mmap = true; |
53 | read_options.create_new = true; |
54 | int ret = delete_store.open(read_options); |
55 | ASSERT_EQ(ret, 0); |
56 | |
57 | IDMap id_map("teachers" , "./teachers/" ); |
58 | ret = id_map.open(read_options); |
59 | ASSERT_EQ(ret, 0); |
60 | |
61 | SegmentMeta segment_meta; |
62 | segment_meta.segment_id = 0; |
63 | |
64 | MemorySegmentPtr memory_segment = |
65 | MemorySegment::Create("teachers" , "./teachers/" , segment_meta, |
66 | schema_.get(), &delete_store, &id_map, 5); |
67 | ASSERT_TRUE(memory_segment != nullptr); |
68 | |
69 | ret = memory_segment->open(read_options); |
70 | ASSERT_EQ(ret, 0); |
71 | |
72 | for (size_t i = 0; i < 1000; i++) { |
73 | Record record; |
74 | record.primary_key = i; |
75 | record.lsn = i; |
76 | record.forward_data = std::string("hello" ) + std::to_string(i); |
77 | |
78 | CollectionDataset::ColumnData new_column; |
79 | new_column.column_name = "face" ; |
80 | new_column.data_type = DataTypes::VECTOR_FP32; |
81 | new_column.dimension = 16U; |
82 | |
83 | std::vector<float> fvec(16U); |
84 | for (size_t j = 0; j < 16U; j++) { |
85 | fvec[j] = i * 1.0f; |
86 | } |
87 | std::string vector((char *)fvec.data(), fvec.size() * sizeof(float)); |
88 | new_column.data = vector; |
89 | |
90 | record.column_datas.emplace_back(new_column); |
91 | |
92 | idx_t doc_id; |
93 | ret = memory_segment->insert(record, &doc_id); |
94 | ASSERT_EQ(ret, 0); |
95 | ASSERT_EQ(doc_id, i); |
96 | |
97 | id_map.insert(record.primary_key, doc_id); |
98 | } |
99 | |
100 | for (size_t i = 0; i < 1000; i++) { |
101 | std::vector<float> fvec(16U); |
102 | for (size_t j = 0; j < 16U; j++) { |
103 | fvec[j] = i * 1.0f; |
104 | } |
105 | std::string query((char *)fvec.data(), fvec.size() * sizeof(float)); |
106 | QueryParams query_params; |
107 | query_params.topk = 10; |
108 | query_params.data_type = DataTypes::VECTOR_FP32; |
109 | query_params.dimension = 16; |
110 | |
111 | QueryResultList result_list; |
112 | ret = memory_segment->knn_search("face" , query, query_params, &result_list); |
113 | ASSERT_EQ(ret, 0); |
114 | |
115 | ASSERT_EQ(result_list[0].primary_key, i); |
116 | ASSERT_EQ(result_list[0].score, 0.0f); |
117 | ASSERT_EQ(result_list[0].lsn, i); |
118 | ASSERT_EQ(result_list[0].forward_data, |
119 | std::string("hello" ) + std::to_string(i)); |
120 | } |
121 | |
122 | for (size_t i = 0; i < 1000; i++) { |
123 | std::vector<float> fvec(16U); |
124 | for (size_t j = 0; j < 16U; j++) { |
125 | fvec[j] = i * 1.0f; |
126 | } |
127 | std::string query((char *)fvec.data(), fvec.size() * sizeof(float)); |
128 | QueryParams query_params; |
129 | query_params.topk = 10; |
130 | query_params.data_type = DataTypes::VECTOR_FP32; |
131 | query_params.dimension = 16; |
132 | query_params.radius = 0.1f; |
133 | |
134 | QueryResultList result_list; |
135 | ret = memory_segment->knn_search("face" , query, query_params, &result_list); |
136 | ASSERT_EQ(ret, 0); |
137 | |
138 | ASSERT_EQ(result_list.size(), 1); |
139 | ASSERT_EQ(result_list[0].primary_key, i); |
140 | ASSERT_EQ(result_list[0].score, 0.0f); |
141 | ASSERT_EQ(result_list[0].lsn, i); |
142 | ASSERT_EQ(result_list[0].forward_data, |
143 | std::string("hello" ) + std::to_string(i)); |
144 | } |
145 | |
146 | for (size_t i = 0; i < 1000; i++) { |
147 | QueryResult result; |
148 | ret = memory_segment->kv_search(i, &result); |
149 | ASSERT_EQ(ret, 0); |
150 | ASSERT_EQ(result.primary_key, i); |
151 | ASSERT_EQ(result.score, 0.0f); |
152 | ASSERT_EQ(result.lsn, i); |
153 | ASSERT_EQ(result.forward_data, std::string("hello" ) + std::to_string(i)); |
154 | } |
155 | } |
156 | |