1 /**
2 * Copyright 2019 Huawei Technologies Co., Ltd
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 <sstream>
18 #include "minddata/dataset/util/btree.h"
19 #include "minddata/dataset/util/auto_index.h"
20 #include "minddata/dataset/util/system_pool.h"
21 #include "minddata/dataset/util/task_manager.h"
22 #include "common/common.h"
23 #include "gtest/gtest.h"
24 #include "utils/log_adapter.h"
25
26 using namespace mindspore::dataset;
27 using mindspore::LogStream;
28 using mindspore::ExceptionType::NoExceptionType;
29 using mindspore::MsLogLevel::INFO;
30
31 // For testing purposes, we will make the branching factor very low.
32 struct mytraits {
33 using slot_type = uint16_t;
34 static const slot_type kLeafSlots = 6;
35 static const slot_type kInnerSlots = 3;
36 };
37
38 class MindDataTestBPlusTree : public UT::Common {
39 public:
40 MindDataTestBPlusTree() = default;
41 };
42
43 // Test serial insert.
TEST_F(MindDataTestBPlusTree,Test1)44 TEST_F(MindDataTestBPlusTree, Test1) {
45 Allocator<std::string> alloc(std::make_shared<SystemPool>());
46 BPlusTree<uint64_t, std::string, Allocator<std::string>, std::less<>, mytraits> btree(alloc);
47 Status rc;
48 for (int i = 0; i < 100; i++) {
49 uint64_t key = 2 * i;
50 std::ostringstream oss;
51 oss << "Hello World. I am " << key;
52 rc = btree.DoInsert(key, oss.str());
53 EXPECT_TRUE(rc.IsOk());
54 }
55 for (int i = 0; i < 100; i++) {
56 uint64_t key = 2 * i + 1;
57 std::ostringstream oss;
58 oss << "Hello World. I am " << key;
59 rc = btree.DoInsert(key, oss.str());
60 EXPECT_TRUE(rc.IsOk());
61 }
62 EXPECT_EQ(btree.size(), 200);
63
64 // Test iterator
65 {
66 int cnt = 0;
67 auto it = btree.begin();
68 uint64_t prev = it.key();
69 ++it;
70 ++cnt;
71 while (it != btree.end()) {
72 uint64_t cur = it.key();
73 std::string val = "Hello World. I am " + std::to_string(cur);
74 EXPECT_TRUE(prev < cur);
75 EXPECT_EQ(it.value(), val);
76 prev = cur;
77 ++it;
78 ++cnt;
79 }
80 EXPECT_EQ(cnt, 200);
81 // Now go backward
82 for (int i = 0; i < 10; i++) {
83 --it;
84 EXPECT_EQ(199 - i, it.key());
85 }
86 }
87
88 // Test search
89 {
90 MS_LOG(INFO) << "Locate key " << 100 << " Expect found.";
91 auto r = btree.Search(100);
92 auto &it = r.first;
93 EXPECT_TRUE(r.second);
94 EXPECT_EQ(it.key(), 100);
95 EXPECT_EQ(it.value(), "Hello World. I am 100");
96 MS_LOG(INFO) << "Locate key " << 300 << " Expect not found.";
97 auto q = btree.Search(300);
98 EXPECT_FALSE(q.second);
99 }
100
101 // Test duplicate key
102 {
103 rc = btree.DoInsert(100, "Expect error");
104 EXPECT_EQ(rc, Status(StatusCode::kMDDuplicateKey));
105 }
106 }
107
108 // Test concurrent insert.
TEST_F(MindDataTestBPlusTree,Test2)109 TEST_F(MindDataTestBPlusTree, Test2) {
110 Allocator<std::string> alloc(std::make_shared<SystemPool>());
111 BPlusTree<uint64_t, std::string, Allocator<std::string>, std::less<>, mytraits> btree(alloc);
112 TaskGroup vg;
113 auto f = [&](int k) -> Status {
114 TaskManager::FindMe()->Post();
115 for (int i = 0; i < 100; i++) {
116 uint64_t key = k * 100 + i;
117 std::ostringstream oss;
118 oss << "Hello World. I am " << key;
119 Status rc = btree.DoInsert(key, oss.str());
120 EXPECT_TRUE(rc.IsOk());
121 }
122 return Status::OK();
123 };
124 auto g = [&](int k) -> Status {
125 TaskManager::FindMe()->Post();
126 for (int i = 0; i < 1000; i++) {
127 uint64_t key = rand() % 10000;
128 ;
129 auto it = btree.Search(key);
130 }
131 return Status::OK();
132 };
133 // Spawn multiple threads to do insert.
134 for (int k = 0; k < 100; k++) {
135 vg.CreateAsyncTask("Concurrent Insert", std::bind(f, k));
136 }
137 // Spawn a few threads to do random search.
138 for (int k = 0; k < 2; k++) {
139 vg.CreateAsyncTask("Concurrent search", std::bind(g, k));
140 }
141 vg.join_all();
142 EXPECT_EQ(btree.size(), 10000);
143
144 // Test iterator
145 {
146 int cnt = 0;
147 auto it = btree.begin();
148 uint64_t prev = it.key();
149 ++it;
150 ++cnt;
151 while (it != btree.end()) {
152 uint64_t cur = it.key();
153 std::string val = "Hello World. I am " + std::to_string(cur);
154 EXPECT_TRUE(prev < cur);
155 EXPECT_EQ(it.value(), val);
156 prev = cur;
157 ++it;
158 ++cnt;
159 }
160 EXPECT_EQ(cnt, 10000);
161 }
162
163 // Test search
164 {
165 MS_LOG(INFO) << "Locating key from 0 to 9999. Expect found.";
166 for (int i = 0; i < 10000; i++) {
167 auto r = btree.Search(i);
168 EXPECT_TRUE(r.second);
169 if (r.second) {
170 auto &it = r.first;
171 EXPECT_EQ(it.key(), i);
172 std::string val = "Hello World. I am " + std::to_string(i);
173 EXPECT_EQ(it.value(), val);
174 }
175 }
176 MS_LOG(INFO) << "Locate key " << 10000 << ". Expect not found";
177 auto q = btree.Search(10000);
178 EXPECT_FALSE(q.second);
179 }
180 }
181
TEST_F(MindDataTestBPlusTree,Test3)182 TEST_F(MindDataTestBPlusTree, Test3) {
183 Allocator<std::string> alloc(std::make_shared<SystemPool>());
184 AutoIndexObj<std::string, Allocator<std::string>> ai(alloc);
185 Status rc;
186 rc = ai.insert("Hello World");
187 EXPECT_TRUE(rc.IsOk());
188 rc = ai.insert({"a", "b", "c"});
189 EXPECT_TRUE(rc.IsOk());
190 uint64_t min = ai.min_key();
191 uint64_t max = ai.max_key();
192 EXPECT_EQ(min, 0);
193 EXPECT_EQ(max, 3);
194 auto r = ai.Search(2);
195 auto &it = r.first;
196 EXPECT_EQ(it.value(), "b");
197 MS_LOG(INFO) << "Dump all the values using [] operator.";
198 for (uint64_t i = min; i <= max; i++) {
199 MS_LOG(DEBUG) << ai[i] << std::endl;
200 }
201 }
202
TEST_F(MindDataTestBPlusTree,Test4)203 TEST_F(MindDataTestBPlusTree, Test4) {
204 Allocator<int64_t> alloc(std::make_shared<SystemPool>());
205 AutoIndexObj<int64_t, Allocator<int64_t>> ai(alloc);
206 Status rc;
207 for (int i = 0; i < 1000; i++) {
208 rc = ai.insert(std::make_unique<int64_t>(i));
209 EXPECT_TRUE(rc.IsOk());
210 }
211 // Test iterator
212 {
213 int cnt = 0;
214 auto it = ai.begin();
215 uint64_t prev = it.key();
216 ++it;
217 ++cnt;
218 while (it != ai.end()) {
219 uint64_t cur = it.key();
220 EXPECT_TRUE(prev < cur);
221 EXPECT_EQ(it.value(), cnt);
222 prev = cur;
223 ++it;
224 ++cnt;
225 }
226 EXPECT_EQ(cnt, 1000);
227 }
228 }
229
TEST_F(MindDataTestBPlusTree,TestPerfNoLocking)230 TEST_F(MindDataTestBPlusTree, TestPerfNoLocking) {
231 AutoIndexObj<int64_t> btree;
232 // No locking test
233 btree.SetLocking(false);
234 // Insert a million entries using the default traits.
235 for (auto i = 0; i < 1000000; ++i) {
236 ASSERT_TRUE(btree.insert(i));
237 }
238 std::cout << "Tree height : " << btree.GetHeight() << std::endl;
239 std::cout << "Tree Order : " << btree.GetOrder() << std::endl;
240 std::cout << "Number of leaves : " << btree.GetNumLeaves() << std::endl;
241 std::cout << "Number of inner nodes : " << btree.GetNumInnerNodes() << std::endl;
242
243 auto r = btree.Search(3);
244 EXPECT_TRUE(r.second);
245 r = btree.Search(999999);
246 EXPECT_TRUE(r.second);
247 }
248