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/index_service.h" |
18 | #include <gtest/gtest.h> |
19 | |
20 | using namespace proxima::be; |
21 | using namespace proxima::be::index; |
22 | |
23 | class IndexServiceTest : public testing::Test { |
24 | protected: |
25 | void SetUp() { |
26 | char cmd_buf[100]; |
27 | snprintf(cmd_buf, 100, "rm -rf ./teachers/" ); |
28 | system(cmd_buf); |
29 | |
30 | snprintf(cmd_buf, 100, "rm -rf ./students/" ); |
31 | system(cmd_buf); |
32 | |
33 | FillSchema(); |
34 | } |
35 | |
36 | void TearDown() {} |
37 | |
38 | void FillSchema() { |
39 | schema_ = std::make_shared<meta::CollectionMeta>(); |
40 | meta::ColumnMetaPtr column_meta = std::make_shared<meta::ColumnMeta>(); |
41 | column_meta->set_name("face" ); |
42 | column_meta->set_index_type(IndexTypes::PROXIMA_GRAPH_INDEX); |
43 | column_meta->set_data_type(DataTypes::VECTOR_FP32); |
44 | column_meta->set_dimension(16); |
45 | column_meta->mutable_parameters()->set("metric_type" , "SquaredEuclidean" ); |
46 | schema_->append(column_meta); |
47 | schema_->set_name("teachers" ); |
48 | |
49 | schema1_ = std::make_shared<meta::CollectionMeta>(); |
50 | meta::ColumnMetaPtr column_meta1 = std::make_shared<meta::ColumnMeta>(); |
51 | column_meta1->set_name("face" ); |
52 | column_meta1->set_index_type(IndexTypes::PROXIMA_GRAPH_INDEX); |
53 | column_meta1->set_data_type(DataTypes::VECTOR_FP32); |
54 | column_meta1->set_dimension(16); |
55 | column_meta1->mutable_parameters()->set("metric_type" , "SquaredEuclidean" ); |
56 | schema1_->append(column_meta1); |
57 | schema1_->set_name("students" ); |
58 | } |
59 | |
60 | meta::CollectionMetaPtr CreateSchema(const std::string &name) { |
61 | auto schema = std::make_shared<meta::CollectionMeta>(); |
62 | meta::ColumnMetaPtr column_meta = std::make_shared<meta::ColumnMeta>(); |
63 | column_meta->set_name("column_test" ); |
64 | column_meta->set_index_type(IndexTypes::PROXIMA_GRAPH_INDEX); |
65 | column_meta->set_data_type(DataTypes::VECTOR_FP32); |
66 | column_meta->set_dimension(16); |
67 | column_meta->mutable_parameters()->set("metric_type" , "SquaredEuclidean" ); |
68 | schema->append(column_meta); |
69 | schema->set_name(name); |
70 | return schema; |
71 | } |
72 | |
73 | protected: |
74 | meta::CollectionMetaPtr schema_{}; |
75 | meta::CollectionMetaPtr schema1_{}; |
76 | }; |
77 | |
78 | TEST_F(IndexServiceTest, TestGeneral) { |
79 | IndexService index_service; |
80 | int ret = index_service.init(); |
81 | ASSERT_EQ(ret, 0); |
82 | |
83 | ret = index_service.start(); |
84 | ASSERT_EQ(ret, 0); |
85 | |
86 | ret = index_service.create_collection("teachers" , schema_); |
87 | ASSERT_EQ(ret, 0); |
88 | |
89 | ASSERT_EQ(index_service.has_collection("teachers" ), true); |
90 | |
91 | std::vector<std::string> names; |
92 | ret = index_service.list_collections(&names); |
93 | ASSERT_EQ(ret, 0); |
94 | ASSERT_EQ(names.size(), 1); |
95 | ASSERT_EQ(names[0], "teachers" ); |
96 | |
97 | std::vector<SegmentPtr> segments; |
98 | ret = index_service.list_segments("teachers" , &segments); |
99 | ASSERT_EQ(ret, 0); |
100 | ASSERT_EQ(segments.size(), 1); |
101 | |
102 | index_service.stop(); |
103 | } |
104 | |
105 | TEST_F(IndexServiceTest, TestLoadCollection) { |
106 | IndexService index_service; |
107 | |
108 | ASSERT_EQ(index_service.init(), 0); |
109 | ASSERT_EQ(index_service.start(), 0); |
110 | |
111 | int ret = index_service.create_collection("teachers" , schema_); |
112 | ASSERT_EQ(ret, 0); |
113 | |
114 | ret = index_service.create_collection("students" , schema1_); |
115 | ASSERT_EQ(ret, 0); |
116 | |
117 | ASSERT_EQ(index_service.stop(), 0); |
118 | |
119 | ret = index_service.start(); |
120 | std::vector<std::string> names; |
121 | names.emplace_back("teachers" ); |
122 | names.emplace_back("students" ); |
123 | |
124 | std::vector<meta::CollectionMetaPtr> schemas; |
125 | schemas.emplace_back(schema_); |
126 | schemas.emplace_back(schema1_); |
127 | |
128 | ret = index_service.load_collections(names, schemas); |
129 | ASSERT_EQ(ret, 0); |
130 | |
131 | index_service.stop(); |
132 | } |
133 | |
134 | void do_hybrid_collection_operations(IndexService *service, |
135 | const std::string &name, |
136 | const meta::CollectionMetaPtr &schema) { |
137 | int ret = service->create_collection(name, schema); |
138 | ASSERT_EQ(ret, 0); |
139 | |
140 | ret = service->drop_collection(name); |
141 | ASSERT_EQ(ret, 0); |
142 | } |
143 | |
144 | TEST_F(IndexServiceTest, TestMultiThread) { |
145 | IndexService index_service; |
146 | |
147 | ASSERT_EQ(index_service.init(), 0); |
148 | ASSERT_EQ(index_service.start(), 0); |
149 | |
150 | ailego::ThreadPool pool(3); |
151 | for (int i = 0; i < 100; i++) { |
152 | std::string name = "collection_" + std::to_string(i); |
153 | auto schema = CreateSchema(name); |
154 | pool.execute(&do_hybrid_collection_operations, &index_service, name, |
155 | schema); |
156 | } |
157 | pool.wait_finish(); |
158 | index_service.stop(); |
159 | } |
160 | |