add: cpp test of top_k data structure.

Author: hidedim

src/types/topk.cc

@@ -76,7 +76,7 @@ static uint32_t MurmurHash2(const void *key, int len, uint32_t seed) {
         len -= 4;
     }
 
-    // Handle the last few bytes of the input array
+    // Handle the last few bytes of the input heap
 
     switch (len) {
     case 3:
@@ -219,65 +219,65 @@ static uint32_t MurmurHash2(const void *key, int len, uint32_t seed) {
 }
 
 /* ---------------------------------------------------------------------- */
-void BlockSplitTopK::swapHeapBucket(HeapBucket *a, HeapBucket* b) {
-    HeapBucket tmp = *a;
-    a->count = b->count;
-    a->fp = b->fp;
-    a->itemlen = b->itemlen;
-    a->item = b->item;
-    
-    b->count = tmp.count;
-    b->fp = tmp.fp;
-    b->itemlen = tmp.itemlen;
-    b->item = tmp.item;
-}
-
 void BlockSplitTopK::heapifyDown(int start) {
-    size_t parent = start;
+    size_t child = start;
 
     // check whether larger than children
-    if (heap_size < 2 || (heap_size - 2) / 2 < parent) {
+    if (heap_size < 2 || (heap_size - 2) / 2 < child) {
         return;
     }
 
-    while (parent < heap_size) {
-        size_t left_child = 2 * parent + 1;
-        size_t right_child = left_child + 1;
+    child = 2 * child + 1;
+    if ((child + 1) < heap_size && (heap[child].count > heap[child + 1].count)) {
+        ++child;
+    }
+    if (heap[child].count > heap[start].count) {
+        return;
+    }
 
-        // check whether left child is larger than parent
-        if (left_child < heap_size && heap[left_child].count < heap[parent].count) {
-            parent = left_child;
-        }
-        // check whether right child is larger than parent
-        if (right_child < heap_size && heap[right_child].count < heap[parent].count) {
-            parent = right_child;
-        }
-        if (parent == left_child || parent == right_child) {
-            swapHeapBucket(&heap[(parent-1)/2], &heap[parent]);
-        } else {
+    HeapBucket top;
+    memcpy(&top, &heap[start], sizeof(HeapBucket));
+    do {
+        memcpy(&heap[start], &heap[child], sizeof(HeapBucket));
+        start = child;
+
+        if ((heap_size - 2) / 2 < child) {
             break;
         }
-    }
+        child = 2 * child + 1;
+
+        if ((child + 1) < heap_size && (heap[child].count > heap[child + 1].count)) {
+            ++child;
+        }
+    } while (heap[child].count < top.count);
+    memcpy(&heap[start], &top, sizeof(HeapBucket));
 }
 
 void BlockSplitTopK::heapifyUp(int start) {
-    size_t child = start;
+    size_t parent = start;
 
     // check whether smaller than parent
-    if (heap_size < 2 || child == 0) {
+    if (heap_size < 2 || parent == 0) {
         return;
     }
 
-    while (child > 0) {
-        size_t parent = (child - 1) / 2;
+    parent = (parent - 1) / 2;
+    if (heap[parent].count > heap[start].count) {
+        return;
+    }
 
-        if (parent >= 0 && heap[child].count < heap[parent].count) {
-            swapHeapBucket(&heap[parent], &heap[child]);
-            child = parent;
-        } else {
+    HeapBucket bottom;
+    memcpy(&bottom, &heap[start], sizeof(HeapBucket));
+    do {
+        memcpy(&heap[start], &heap[parent], sizeof(HeapBucket));
+        start = parent;
+
+        if (start == 0) {
             break;
         }
-    }
+        parent = (parent - 1) / 2;
+    } while (heap[parent].count > bottom.count);
+    memcpy(&heap[start], &bottom, sizeof(HeapBucket));
 }
 
 int BlockSplitTopK::checkExistInHeap(const std::string &item) {
@@ -310,8 +310,10 @@ void BlockSplitTopK::Add(const std::string &item, uint32_t increment) {
         if (buckets[loc].count == 0) {
             buckets[loc].fp = fp;
             buckets[loc].count = increment;
+            maxCount = std::max(maxCount, buckets[loc].count);
         } else if (buckets[loc].fp == fp && location != -1) {
             buckets[loc].count += increment;
+            maxCount = std::max(maxCount, buckets[loc].count);
         } else {
             // decay
             uint32_t local_incr = increment;
@@ -330,12 +332,12 @@ void BlockSplitTopK::Add(const std::string &item, uint32_t increment) {
                     if (buckets[loc].count == 0) {
                         buckets[loc].fp = fp;
                         buckets[loc].count = 1;
+                        maxCount = std::max(maxCount, buckets[loc].count);
                         break;
                     }
                 }
             }
         }
-        maxCount = std::max(maxCount, buckets[loc].count);
     }
 
     if (k == heap_size) {

src/types/topk.h

@@ -72,7 +72,6 @@ class BlockSplitTopK {
     void heapifyUp(int start);
     int checkExistInHeap(const std::string &item);
     int cmpHeapBucketCount(const HeapBucket &a, const HeapBucket &b);
-    void swapHeapBucket(HeapBucket *a, HeapBucket *b);
 
     uint32_t k;
     uint32_t width;

tests/cppunit/types/topk_test.cc

@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#include <gtest/gtest.h>
+
+#include <memory>
+
+#include "test_base.h"
+#include "types/redis_topk.h"
+
+static constexpr uint32_t k = 5;
+static constexpr uint32_t width = 7;
+static constexpr uint32_t depth = 8;
+static constexpr double decay = 0.9;
+
+class RedisTopKTest : public TestBase {
+ protected:
+  explicit RedisTopKTest() : TestBase() {
+    top_k_ = std::make_unique<redis::TopK>(storage_.get(), "topk_ns");
+  }
+  ~RedisTopKTest() override = default;
+
+  void SetUp() override {
+    key_ = "test_topk->key";
+    top_k_->Reserve(*ctx_, key_, k, width, depth, decay);
+  }
+
+  void TearDown() override {}
+
+  std::unique_ptr<redis::TopK> top_k_;
+};
+
+TEST_F(RedisTopKTest, TestTopKInfo) {
+  // test exist key
+  redis::TopKInfo info1;
+  top_k_->Info(*ctx_, key_, &info1);
+  ASSERT_EQ(info1.k, k);
+  ASSERT_EQ(info1.width, width);
+  ASSERT_EQ(info1.depth, depth);
+  ASSERT_EQ(info1.decay, decay);
+
+  // test not exist key
+  redis::TopKInfo info2;
+  auto s = top_k_->Info(*ctx_, "not_exist_key", &info2);
+  ASSERT_FALSE(s.ok());
+}
+
+TEST_F(RedisTopKTest, TestTopKAddAndQuery) {
+  // test not exist key
+  std::string no_exist_key = "no_exist_key";
+  auto s = top_k_->Add(*ctx_, no_exist_key, "1");
+  ASSERT_FALSE(s.ok());
+
+  bool exist;
+  s = top_k_->Query(*ctx_, no_exist_key, "1", &exist);
+  ASSERT_FALSE(s.ok());
+
+  std::vector<std::string> list;
+  s = top_k_->List(*ctx_, no_exist_key, list);
+  ASSERT_FALSE(s.ok());
+
+  // test exist key
+  std::vector<std::string> values1 = {"1", "2", "3", "4", "5"};
+  std::vector<std::string> values2 = {"6", "7", "8", "9", "10"};
+  std::unordered_set<std::string> values_set1(values1.begin(), values1.end());
+  std::unordered_set<std::string> values_set2(values2.begin(), values2.end());
+
+  // found not exist values1
+  for (size_t i = 0; i < values1.size(); ++i) {
+    bool found = true;
+    top_k_->Query(*ctx_, key_, values1[i], &found);
+    ASSERT_FALSE(found);
+  }
+  // add values1, and query values1.
+  for (size_t i = 0; i < values1.size(); ++i) {
+    top_k_->Add(*ctx_, key_, values1[i]);
+    bool found = false;
+    top_k_->Query(*ctx_, key_, values1[i], &found);
+    ASSERT_TRUE(found);
+  }
+  for (size_t i = 0; i < values1.size(); ++i) {
+    bool found = false;
+    top_k_->Query(*ctx_, key_, values1[i], &found);
+    ASSERT_TRUE(found);
+  }
+
+  // found topk list.
+  std::vector<std::string> top_k_list;
+  top_k_->List(*ctx_, key_, top_k_list);
+  ASSERT_EQ(top_k_list.size(), values1.size());
+  for (size_t i = 0; i < k; ++i) {
+    ASSERT_TRUE(values_set1.find(top_k_list[i]) != values_set1.end());
+  }
+
+  // heap is full, need remove values1.
+  for (size_t i = 0; i < values2.size(); ++i) {
+    bool found = false;
+    // due to decay, topk is possiable to remove values1.
+    while (!found) {
+      top_k_->Add(*ctx_, key_, values2[i]);
+      top_k_->Query(*ctx_, key_, values2[i], &found);
+    }
+    top_k_->Add(*ctx_, key_, values2[i]);
+  }
+
+  // values1 is removed.
+  for (size_t i = 0; i < values1.size(); ++i) {
+    bool found = true;
+    top_k_->Query(*ctx_, key_, values1[i], &found);
+    ASSERT_FALSE(found);
+  }
+
+  // found topk list.
+  top_k_list.clear();
+  top_k_->List(*ctx_, key_, top_k_list);
+  for (size_t i = 0; i < top_k_list.size(); ++i) {
+    ASSERT_TRUE(values_set2.find(top_k_list[i]) != values_set2.end());
+  }
+}