chore: optimize zpopminmax operations (#3291)

* chore: optimize zpopminmax operations

Before, we run FromRank operation count times. This operation computes subtree counts on the fly.
Now, we iterate over a tree first, delete the keys from the map - no rank calls.
And only after we filled the result, we delete the entries from the tree.
For cases where we delete all the entries we avoid calling FromRank.

Moreover, we optimized FromRank operations to handle corner cases of most right and most left nodes
by avoiding computing the subtree counts.

---------

Signed-off-by: Roman Gershman <[email protected]>

src/core/bptree_set.h

@@ -79,12 +79,14 @@ template <typename T, typename Policy = BPTreePolicy<T>> class BPTree {
   /// @param rank_start
   /// @param rank_end - inclusive.
   /// @param cb - callback to be called for each item in the range.
+  ///             Should return false to stop iteration.
   bool Iterate(uint32_t rank_start, uint32_t rank_end, std::function<bool(KeyT)> cb) const;
 
   /// @brief Iterates over all items in the range [rank_start, rank_end] by rank in reverse order.
   /// @param rank_start
   /// @param rank_end
-  /// @param cb
+  /// @param cb - callback to be called for each item in the range.
+  ///             Should return false to stop iteration.
   bool IterateReverse(uint32_t rank_start, uint32_t rank_end, std::function<bool(KeyT)> cb) const;
 
   /// @brief Returns the path to the first item in the tree that is greater or equal to key.
@@ -435,7 +437,24 @@ void BPTree<T, Policy>::ToRank(uint32_t rank, BPTreePath* path) const {
   assert(root_ && rank < count_);
   BPTreeNode* node = root_;
 
+  if (rank + 1 == count_) {
+    // Corner case where we search for the node on the right.
+    while (!node->IsLeaf()) {
+      path->Push(node, node->NumItems());
+      node = node->Child(node->NumItems());
+    }
+    path->Push(node, node->NumItems() - 1);
+    return;
+  }
+
   while (!node->IsLeaf()) {
+    // handle common corner case of search of left-most node, and avoid counting sub-tree count.
+    if (rank == 0) {
+      path->Push(node, 0);
+      node = node->Child(0);
+      continue;
+    }
+
     for (unsigned i = 0; i <= node->NumItems(); ++i) {
       uint32_t subtree_cnt = node->GetChildTreeCount(i);
       if (subtree_cnt > rank) {

src/core/sorted_map.cc

@@ -573,34 +573,50 @@ size_t SortedMap::DeleteRangeByLex(const zlexrangespec& range) {
 }
 
 SortedMap::ScoredArray SortedMap::PopTopScores(unsigned count, bool reverse) {
+  DCHECK_GT(count, 0u);
   DCHECK_EQ(score_map->UpperBoundSize(), score_tree->Size());
   size_t sz = score_map->UpperBoundSize();
 
   ScoredArray res;
 
-  if (sz == 0)
+  DCHECK_GT(sz, 0u);  // Empty sets are not allowed.
+
+  if (sz == 0 || count == 0)
     return res;
 
-  if (count >= sz)
-    count = score_map->UpperBoundSize();
+  if (count > sz)
+    count = sz;
 
   res.reserve(count);
+
+  auto cb = [&](ScoreSds obj) {
+    res.emplace_back(string{(sds)obj, sdslen((sds)obj)}, GetObjScore(obj));
+
+    // We can not delete from score_tree because we are in the middle of the iteration.
+    CHECK(score_map->Erase((sds)obj));
+    return true;  // continue with the iteration.
+  };
+
   unsigned rank = 0;
   unsigned step = 0;
-
   if (reverse) {
-    rank = sz - 1;
+    score_tree->IterateReverse(0, count - 1, std::move(cb));
+    rank = score_tree->Size() - 1;
     step = 1;
+  } else {
+    score_tree->Iterate(0, count - 1, std::move(cb));
   }
 
-  for (unsigned i = 0; i < count; ++i) {
-    auto path = score_tree->FromRank(rank);
-    ScoreSds obj = path.Terminal();
-    res.emplace_back(string{(sds)obj, sdslen((sds)obj)}, GetObjScore(obj));
-
-    score_tree->Delete(path);
-    score_map->Erase((sds)obj);
-    rank -= step;
+  // We already deleted elements from score_map, so what's left is to delete from the tree.
+  if (score_map->Empty()) {
+    // Corner case optimization.
+    score_tree->Clear();
+  } else {
+    for (unsigned i = 0; i < res.size(); ++i) {
+      auto path = score_tree->FromRank(rank);
+      score_tree->Delete(path);
+      rank -= step;
+    }
   }
 
   return res;

src/core/sorted_map.h

@@ -54,6 +54,9 @@ class SortedMap {
   bool Insert(double score, sds member);
   bool Delete(sds ele);
 
+  // Upper bound size of the set.
+  // Note: Currently we do not allow member expiry in sorted sets, therefore it's exact
+  // But if we decide to add expire, this method will provide an approximation from above.
   size_t Size() const {
     return score_map->UpperBoundSize();
   }
@@ -87,8 +90,11 @@ class SortedMap {
  private:
   using ScoreTree = BPTree<ScoreSds, ScoreSdsPolicy>;
 
+  // hash map from fields to scores.
   ScoreMap* score_map = nullptr;
-  ScoreTree* score_tree = nullptr;  // just a stub for now.
+
+  // sorted tree of (score,field) items.
+  ScoreTree* score_tree = nullptr;
 };
 
 // Used by CompactObject.

src/server/zset_family.cc

@@ -487,11 +487,13 @@ void IntervalVisitor::ActionRem(const zlexrangespec& range) {
 }
 
 void IntervalVisitor::ActionPop(ZSetFamily::TopNScored sc) {
-  if (robj_wrapper_->encoding() == OBJ_ENCODING_LISTPACK) {
-    PopListPack(sc);
-  } else {
-    CHECK_EQ(robj_wrapper_->encoding(), OBJ_ENCODING_SKIPLIST);
-    PopSkipList(sc);
+  if (sc > 0) {
+    if (robj_wrapper_->encoding() == OBJ_ENCODING_LISTPACK) {
+      PopListPack(sc);
+    } else {
+      CHECK_EQ(robj_wrapper_->encoding(), OBJ_ENCODING_SKIPLIST);
+      PopSkipList(sc);
+    }
   }
 }
 

src/server/zset_family_test.cc

@@ -698,6 +698,9 @@ TEST_F(ZSetFamilyTest, ZPopMin) {
   ASSERT_THAT(resp, ArrLen(2));
   EXPECT_THAT(resp.GetVec(), ElementsAre("a", "1"));
 
+  resp = Run({"zpopmin", "key", "0"});
+  ASSERT_THAT(resp, ArrLen(0));
+
   resp = Run({"zpopmin", "key", "2"});
   ASSERT_THAT(resp, ArrLen(4));
   EXPECT_THAT(resp.GetVec(), ElementsAre("b", "2", "c", "3"));