add zpopmax zpopmin (OpenAtomFoundation#706)

Atlas-2020 · Aug 7, 2019 · 9abee07 · 9abee07
1 parent 059c8e6
commit 9abee07
Show file tree

Hide file tree

Showing 5 changed files with 112 additions and 1 deletion.
diff --git a/include/pika_command.h b/include/pika_command.h
@@ -143,6 +143,8 @@ const std::string kCmdNameZLexcount = "zlexcount";
 const std::string kCmdNameZRemrangebyrank = "zremrangebyrank";
 const std::string kCmdNameZRemrangebylex = "zremrangebylex";
 const std::string kCmdNameZRemrangebyscore = "zremrangebyscore";
+const std::string kCmdNameZPopmax = "zpopmax";
+const std::string kCmdNameZPopmin = "zpopmin";
 
 //Set
 const std::string kCmdNameSAdd = "sadd";

diff --git a/include/pika_zset.h b/include/pika_zset.h
@@ -411,4 +411,37 @@ class ZRemrangebylexCmd : public Cmd {
     left_close_ =  right_close_ = true;
   }
 };
+
+class ZPopmaxCmd : public Cmd {
+ public:
+  ZPopmaxCmd(const std::string& name, int arity, uint16_t flag)
+       : Cmd(name, arity, flag) {}
+  virtual std::vector<std::string> current_key() const { 
+    std::vector<std::string> res;
+    res.emplace_back(key_);
+    return res; 
+  }
+  virtual void Do(std::shared_ptr<Partition> partition = nullptr);
+ private:
+  virtual void DoInitial() override;
+  std::string key_;
+  int64_t count_;
+};
+
+class ZPopminCmd : public Cmd {
+ public:
+  ZPopminCmd(const std::string& name, int arity, uint16_t flag)
+       : Cmd(name, arity, flag) {}
+  virtual std::vector<std::string> current_key() const { 
+    std::vector<std::string> res;
+    res.push_back(key_);
+    return res; 
+  }
+  virtual void Do(std::shared_ptr<Partition> partition = nullptr);
+ private:
+  virtual void DoInitial() override;
+  std::string key_;
+  int64_t count_;
+};
+
 #endif
diff --git a/src/pika_command.cc b/src/pika_command.cc
@@ -369,6 +369,12 @@ void InitCmdTable(std::unordered_map<std::string, Cmd*> *cmd_table) {
   ////ZRemrangebylexCmd
   Cmd* zremrangebylexptr = new ZRemrangebylexCmd(kCmdNameZRemrangebylex, 4, kCmdFlagsWrite | kCmdFlagsSinglePartition | kCmdFlagsZset);
   cmd_table->insert(std::pair<std::string, Cmd*>(kCmdNameZRemrangebylex, zremrangebylexptr));
+  ////ZPopmax
+  Cmd* zpopmax = new ZPopmaxCmd(kCmdNameZPopmax, -2, kCmdFlagsWrite | kCmdFlagsSinglePartition | kCmdFlagsZset);
+  cmd_table->insert(std::pair<std::string, Cmd*>(kCmdNameZPopmax, zpopmax));
+  ////ZPopmin
+  Cmd* zpopmin = new ZPopminCmd(kCmdNameZPopmin, -2, kCmdFlagsWrite | kCmdFlagsSinglePartition | kCmdFlagsZset);
+  cmd_table->insert(std::pair<std::string, Cmd*>(kCmdNameZPopmin, zpopmin));
 
   //Set
   ////SAddCmd

diff --git a/src/pika_zset.cc b/src/pika_zset.cc
@@ -874,3 +874,73 @@ void ZRemrangebylexCmd::Do(std::shared_ptr<Partition> partition) {
   res_.AppendInteger(count);
   return;
 }
+
+
+void ZPopmaxCmd::DoInitial() {
+  if (!CheckArg(argv_.size())) {
+    res_.SetRes(CmdRes::kWrongNum, kCmdNameZPopmax);
+    return;
+  }
+  key_ = argv_[1];
+  if (argv_.size() == 2) {
+    count_ = 1;
+    return;
+  }  
+  if (!slash::string2ll(argv_[2].data(), argv_[2].size(), (long long*)(&count_))) {
+    res_.SetRes(CmdRes::kInvalidInt);
+    return;
+  }
+}
+
+void ZPopmaxCmd::Do(std::shared_ptr<Partition> partition) {
+  std::vector<blackwidow::ScoreMember> score_members;
+  rocksdb::Status s = partition->db()->ZPopMax(key_, count_, &score_members);
+  if (s.ok()) {
+      char buf[32];
+      int64_t len;
+      res_.AppendArrayLen(score_members.size() * 2);
+      for (const auto& sm : score_members) {
+        res_.AppendString(sm.member);
+        len = slash::d2string(buf, sizeof(buf), sm.score);
+        res_.AppendStringLen(len);
+        res_.AppendContent(buf);
+      }
+  } else {
+    res_.SetRes(CmdRes::kErrOther, s.ToString());
+  }
+}
+
+
+void ZPopminCmd::DoInitial() {
+  if (!CheckArg(argv_.size())) {
+    res_.SetRes(CmdRes::kWrongNum, kCmdNameZPopmin);
+    return;
+  }
+  key_ = argv_[1];
+  if (argv_.size() == 2) {
+    count_ = 1;
+    return;
+  }  
+  if (!slash::string2ll(argv_[2].data(), argv_[2].size(), (long long*)(&count_))) {
+    res_.SetRes(CmdRes::kInvalidInt);
+    return;
+  }
+}
+
+void ZPopminCmd::Do(std::shared_ptr<Partition> partition) {
+  std::vector<blackwidow::ScoreMember> score_members;
+  rocksdb::Status s = partition->db()->ZPopMin(key_, count_, &score_members);
+  if (s.ok()) {
+      char buf[32];
+      int64_t len;
+      res_.AppendArrayLen(score_members.size() * 2);
+      for (const auto& sm : score_members) {
+        res_.AppendString(sm.member);
+        len = slash::d2string(buf, sizeof(buf), sm.score);
+        res_.AppendStringLen(len);
+        res_.AppendContent(buf);
+      }
+  } else {
+    res_.SetRes(CmdRes::kErrOther, s.ToString());
+  }
+}
diff --git a/third/blackwidow b/third/blackwidow
+18 −1		include/blackwidow/blackwidow.h
+12 −0		src/blackwidow.cc
+96 −0		src/redis_zsets.cc
+6 −0		src/redis_zsets.h
+318 −0		tests/gtest_zsets.cc