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 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_AUTO_INDEX_H_ 17 #define MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_AUTO_INDEX_H_ 18 19 #include <atomic> 20 #include <functional> 21 #include <memory> 22 #include <utility> 23 #include <vector> 24 25 #include "minddata/dataset/util/btree.h" 26 #include "minddata/dataset/util/system_pool.h" 27 28 namespace mindspore { 29 namespace dataset { 30 /// This is a B+ tree with generated int64_t value as key. 31 /// Use minKey() function to query the min key. 32 /// Use maxKey() function to query the max key. 33 /// @tparam T 34 template <typename V, typename A = std::allocator<V>, typename T = BPlusTreeTraits> 35 class AutoIndexObj : public BPlusTree<int64_t, V, A, std::less<int64_t>, T> { 36 public: 37 using my_tree = BPlusTree<int64_t, V, A, std::less<int64_t>, T>; 38 using key_type = typename my_tree::key_type; 39 using value_type = typename my_tree::value_type; 40 AutoIndexObj()41 AutoIndexObj() : my_tree::BPlusTree(), inx_(kMinKey) {} 42 AutoIndexObj(const Allocator<V> & alloc)43 explicit AutoIndexObj(const Allocator<V> &alloc) : my_tree::BPlusTree(alloc), inx_(kMinKey) {} 44 45 ~AutoIndexObj() = default; 46 47 // Insert an object into the tree. 48 // @param val 49 // @return 50 Status insert(const value_type &val, key_type *key = nullptr) { 51 key_type my_inx = inx_.fetch_add(1); 52 if (key != nullptr) { 53 *key = my_inx; 54 } 55 return my_tree::DoInsert(my_inx, val); 56 } 57 58 Status insert(std::unique_ptr<value_type> &&val, key_type *key = nullptr) { 59 key_type my_inx = inx_.fetch_add(1); 60 if (key != nullptr) { 61 *key = my_inx; 62 } 63 return my_tree::DoInsert(my_inx, std::move(val)); 64 } 65 66 // Insert a vector of objects into the tree. 67 // @param v 68 // @return insert(std::vector<value_type> v)69 Status insert(std::vector<value_type> v) { 70 uint64_t num_ele = v.size(); 71 if (num_ele > 0) { 72 // reserve a range of keys rather than getting it one by one. 73 key_type my_inx = inx_.fetch_add(num_ele); 74 for (uint64_t i = 0; i < num_ele; i++) { 75 RETURN_IF_NOT_OK(my_tree::DoInsert(my_inx + i, v.at(i))); 76 } 77 } 78 return Status::OK(); 79 } 80 81 // @return the minimum key min_key()82 key_type min_key() const { 83 auto it = this->cbegin(); 84 return it.key(); 85 } 86 87 // @return the maximum key max_key()88 key_type max_key() const { 89 auto it = this->cend(); 90 --it; 91 return it.key(); 92 } 93 94 private: 95 static constexpr key_type kMinKey = 0; 96 std::atomic<key_type> inx_; 97 }; 98 } // namespace dataset 99 } // namespace mindspore 100 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_AUTO_INDEX_H_ 101