fix: fix prefix lookup in basickv and add corresponding test

zz-jason · Sep 14, 2024 · 92a561a · 92a561a
1 parent a19fa4f
commit 92a561a
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 33 deletions.
diff --git a/src/btree/BasicKV.cpp b/src/btree/BasicKV.cpp
@@ -11,6 +11,9 @@
 #include "leanstore/utils/Misc.hpp"
 
 #include <format>
+#include <string>
+
+#include <sys/types.h>
 
 using namespace std;
 using namespace leanstore::storage;
@@ -178,42 +181,31 @@ OpCode BasicKV::Insert(Slice key, Slice val) {
   return OpCode::kOK;
 }
 
-OpCode BasicKV::PrefixLookup(Slice key, PrefixLookupCallback callback) {
-  while (true) {
-    JUMPMU_TRY() {
-      GuardedBufferFrame<BTreeNode> guardedLeaf;
-      FindLeafCanJump(key, guardedLeaf);
-
-      bool isEqual = false;
-      int16_t cur = guardedLeaf->LowerBound<false>(key, &isEqual);
-      if (isEqual) {
-        callback(key, guardedLeaf->Value(cur));
-        guardedLeaf.JumpIfModifiedByOthers();
-        JUMPMU_RETURN OpCode::kOK;
-      }
-
-      if (cur < guardedLeaf->mNumSlots) {
-        auto fullKeySize = guardedLeaf->GetFullKeyLen(cur);
-        auto fullKeyBuf = utils::JumpScopedArray<uint8_t>(fullKeySize);
-        guardedLeaf->CopyFullKey(cur, fullKeyBuf->get());
-        guardedLeaf.JumpIfModifiedByOthers();
-
-        callback(Slice(fullKeyBuf->get(), fullKeySize), guardedLeaf->Value(cur));
-        guardedLeaf.JumpIfModifiedByOthers();
-
-        JUMPMU_RETURN OpCode::kOK;
+OpCode BasicKV::PrefixLookup(Slice prefixKey, PrefixLookupCallback callback) {
+  JUMPMU_TRY() {
+    auto iter = GetIterator();
+    if (iter.SeekToFirstGreaterEqual(prefixKey); !iter.Valid()) {
+      JUMPMU_RETURN OpCode::kNotFound;
+    }
+    bool foundPrefixKey = false;
+    uint16_t prefixSize = prefixKey.size();
+    for (; iter.Valid(); iter.Next()) {
+      iter.AssembleKey();
+      auto key = iter.Key();
+      auto value = iter.Val();
+      if ((key.size() < prefixSize) || (bcmp(key.data(), prefixKey.data(), prefixSize) != 0)) {
+        break;
       }
-
-      OpCode ret = ScanAsc(key, [&](Slice scannedKey, Slice scannedVal) {
-        callback(scannedKey, scannedVal);
-        return false;
-      });
-      JUMPMU_RETURN ret;
+      callback(key, value);
+      foundPrefixKey = true;
     }
-    JUMPMU_CATCH() {
+    if (!foundPrefixKey) {
+      JUMPMU_RETURN OpCode::kNotFound;
     }
+    JUMPMU_RETURN OpCode::kOK;
+  }
+  JUMPMU_CATCH() {
   }
-
   UNREACHABLE();
   return OpCode::kOther;
 }

diff --git a/tests/btree/BasicKVTest.cpp b/tests/btree/BasicKVTest.cpp
@@ -299,7 +299,7 @@ TEST_F(BasicKVTest, SameKeyInsertRemoveMultiTimes) {
     btree = res.value();
   });
 
-  // insert 100 key-values to the btree
+  // insert 1000 key-values to the btree
   size_t numKVs(1000);
   std::vector<std::pair<std::string, std::string>> kvToTest;
   for (size_t i = 0; i < numKVs; ++i) {
@@ -342,4 +342,79 @@ TEST_F(BasicKVTest, SameKeyInsertRemoveMultiTimes) {
   mStore->ExecSync(0, [&]() { EXPECT_EQ(btree->CountEntries(), numKVs); });
 }
 
+TEST_F(BasicKVTest, PrefixLookup) {
+  storage::btree::BasicKV* btree;
+  mStore->ExecSync(0, [&]() {
+    auto res = mStore->CreateBasicKV(genBtreeName("_tree1"));
+    ASSERT_TRUE(res);
+    ASSERT_NE(res.value(), nullptr);
+    btree = res.value();
+  });
+  // callback function
+  std::vector<std::tuple<std::string, std::string>> copiedKeyValue;
+  auto copyKeyAndValueOut = [&](Slice key, Slice val) {
+    copiedKeyValue.emplace_back(key.ToString(), val.ToString());
+  };
+
+  {
+    // not found valid prefix key
+    std::string prefixString("key_");
+    auto prefixKey = Slice(prefixString);
+    EXPECT_EQ(btree->PrefixLookup(prefixKey, copyKeyAndValueOut), OpCode::kNotFound);
+  }
+
+  {
+    // insert key and value
+    size_t numKVs(10);
+    std::vector<std::pair<std::string, std::string>> kvToTest;
+    for (size_t i = 0; i < numKVs; ++i) {
+      std::string key("key_" + std::string(10, 'x') + std::to_string(i));
+      std::string val("val_" + std::string(100, 'x') + std::to_string(i));
+      mStore->ExecSync(0, [&]() { EXPECT_EQ(btree->Insert(key, val), OpCode::kOK); });
+      kvToTest.emplace_back(std::move(key), std::move(val));
+    }
+
+    // prefix lookup the full set
+    auto prefixString("key_" + std::string(10, 'x'));
+    auto prefixKey = Slice(prefixString);
+    EXPECT_EQ(btree->PrefixLookup(prefixKey, copyKeyAndValueOut), OpCode::kOK);
+    EXPECT_EQ(copiedKeyValue.size(), kvToTest.size());
+    for (size_t i = 0; i < copiedKeyValue.size(); i++) {
+      const auto& [key, expectedVal] = kvToTest[i];
+      const auto& [copiedKey, copiedValue] = copiedKeyValue[i];
+      EXPECT_EQ(copiedKey, key);
+      EXPECT_EQ(copiedValue, expectedVal);
+    }
+
+    // insert special key for prefix lookup
+    std::vector<std::pair<std::string, std::string>> kvToTest2;
+    for (size_t i = 0; i < numKVs; ++i) {
+      std::string key("prefix_key_" + std::string(10, 'x') + std::to_string(i));
+      std::string val("prefix_value_" + std::string(100, 'x') + std::to_string(i));
+      mStore->ExecSync(0, [&]() { EXPECT_EQ(btree->Insert(key, val), OpCode::kOK); });
+      kvToTest2.emplace_back(std::move(key), std::move(val));
+    }
+
+    // prefix lookup the partial set
+    copiedKeyValue.clear();
+    auto prefixString2("prefix_key_" + std::string(10, 'x'));
+    auto prefixKey2 = Slice(prefixString2);
+    EXPECT_EQ(btree->PrefixLookup(prefixKey2, copyKeyAndValueOut), OpCode::kOK);
+    EXPECT_EQ(copiedKeyValue.size(), kvToTest2.size());
+    for (size_t i = 0; i < copiedKeyValue.size(); i++) {
+      const auto& [key, expectedVal] = kvToTest2[i];
+      const auto& [copiedKey, copiedValue] = copiedKeyValue[i];
+      EXPECT_EQ(copiedKey, key);
+      EXPECT_EQ(copiedValue, expectedVal);
+    }
+  }
+  {
+    // greater than the prefix key, but not a true prefix
+    copiedKeyValue.clear();
+    auto prefixString("prefix_kex_" + std::string(10, 'w'));
+    auto prefixKey = Slice(prefixString);
+    EXPECT_EQ(btree->PrefixLookup(prefixKey, copyKeyAndValueOut), OpCode::kNotFound);
+  }
+}
+
 } // namespace leanstore::test