fix(bptree): Reverse iteration (#2075)

Co-authored-by: Roman Gershman <roman@dragonflydb.io>
This commit is contained in:
Shahar Mike 2023-10-25 23:07:57 +03:00 committed by GitHub
parent b2e989a4fb
commit 424c65b4e5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 44 additions and 26 deletions

View file

@ -420,7 +420,7 @@ bool BPTree<T, Policy>::IterateReverse(uint32_t rank_start, uint32_t rank_end,
assert(rank_start <= rank_end && rank_end < count_);
BPTreePath path;
ToRank(rank_end, &path);
ToRank(count_ - 1 - rank_start, &path);
for (uint32_t i = rank_start; i <= rank_end; ++i) {
if (!cb(path.Terminal()))
return false;

View file

@ -4,6 +4,7 @@
#include "core/bptree_set.h"
#include <absl/container/btree_set.h>
#include <gmock/gmock.h>
#include <mimalloc.h>
#include <random>
@ -248,31 +249,6 @@ TEST_F(BPTreeSetTest, Iterate) {
ASSERT_EQ(to - from + 1, cnt);
ASSERT_TRUE(res);
}
// Reverse iteration
cnt = 0;
res = bptree_.IterateReverse(5845, 6849, [&](uint64_t val) {
if ((6849 - cnt) * 2 != val)
return false;
++cnt;
return true;
});
ASSERT_EQ(6849 - 5845 + 1, cnt);
ASSERT_TRUE(res);
for (unsigned j = 0; j < 10; ++j) {
cnt = 0;
unsigned from = generator_() % kNumElems;
unsigned to = from + generator_() % (kNumElems - from);
res = bptree_.IterateReverse(from, to, [&](uint64_t val) {
if ((to - cnt) * 2 != val)
return false;
++cnt;
return true;
});
ASSERT_EQ(to - from + 1, cnt);
ASSERT_TRUE(res);
}
}
TEST_F(BPTreeSetTest, Ranges) {
@ -352,6 +328,48 @@ TEST_F(BPTreeSetTest, InsertSDS) {
}
}
TEST_F(BPTreeSetTest, ReverseIterate) {
vector<ZsetPolicy::KeyT> vals;
for (int i = -1000; i < 1000; ++i) {
sds s = sdsempty();
s = sdscatfmt(s, "a%u", i);
vals.emplace_back(ZsetPolicy::KeyT{.d = (double)i, .s = s});
}
SDSTree tree(&mi_alloc_);
for (auto v : vals) {
ASSERT_TRUE(tree.Insert(v));
{
double score = 0;
tree.IterateReverse(0, 0, [&score](auto i) {
score = i.d;
return false;
});
EXPECT_EQ(score, v.d);
}
{
double score = 0;
tree.Iterate(0, 0, [&score](auto i) {
score = i.d;
return false;
});
EXPECT_EQ(score, vals[0].d);
}
}
vector<int> res;
tree.IterateReverse(0, 1, [&](auto i) {
res.push_back(i.d);
return true;
});
EXPECT_THAT(res, testing::ElementsAre(999, 998));
for (auto v : vals) {
sdsfree(v.s);
}
}
static string RandomString(mt19937& rand, unsigned len) {
const string_view alpanum = "1234567890abcdefghijklmnopqrstuvwxyz";
string ret;