Commit 042781f2 authored by Tat Dat Duong's avatar Tat Dat Duong
Browse files

chore: improve query perf of b-link-tree

parent 5f4d7b33
Loading
Loading
Loading
Loading
+8 −11
Original line number Diff line number Diff line
@@ -256,22 +256,19 @@ public:

  template <typename Warp>
  __device__ static inline bool find(BNode *root, KeyType key, ValueType &result, Warp &warp) {
    auto rank = warp.thread_rank();
    uint64_t rank = warp.thread_rank();

    BNode *next = root;
    while (next) {
      bool isLeaf = next->mLeaf;
      size_t size = next->mSize;

      KeyType nextKey = next->mKeys[rank];
      uint32_t ballot = warp.ballot(nextKey >= key);
      size_t targetIdx = __ffs(ballot) - 1;
      uint16_t isLeaf = next->mLeaf;

      targetIdx = (targetIdx > size) ? size : targetIdx;
      KeyType targetKey = warp.shfl(nextKey, targetIdx);
      uint32_t size = next->mSize;
      uint32_t ballot = warp.ballot(next->mKeys[rank] >= key);
      uint32_t targetIdx = min(__ffs(ballot) - 1, size);
      uint32_t foundKey = next->mKeys[targetIdx] == key;

      if (isLeaf) {
        if (targetIdx < size && targetKey == key) {
        if (foundKey && targetIdx < size) {
          result = next->mValues[targetIdx];
          return true;
        }
@@ -279,7 +276,7 @@ public:
      }

      // assumption: the tree is fully constructed and insert / splits are done
      next = next->mChildren[targetIdx + (targetKey == key)];
      next = next->mChildren[targetIdx + foundKey];
    }

    return false;