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
20using namespace proxima::be;
21using namespace proxima::be::index;
22
23class 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
78TEST_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
105TEST_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
134void 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
144TEST_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