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
21using namespace proxima::be;
22using namespace proxima::be::index;
23
24class 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
49TEST_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