[GitHub] [incubator-kvrocks] torwig commented on a diff in pull request #1502: Add support of the new command ZUNION

["torwig (via GitHub)" <gi...@apache.org>]

torwig commented on code in PR #1502:
URL: https://github.com/apache/incubator-kvrocks/pull/1502#discussion_r1236467536


##########
src/commands/cmd_zset.cc:
##########
@@ -1148,6 +1148,78 @@ class CommandZMScore : public Commander {
   }
 };
 
+class CommandZUnion : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    CommandParser parser(args, 1);
+    numkeys_ = GET_OR_RET(parser.TakeInt<int>(NumericRange<int>{1, std::numeric_limits<int>::max()}));
+    for (size_t i = 0; i < numkeys_; ++i) {
+      keys_weights_.emplace_back(KeyWeight{GET_OR_RET(parser.TakeStr()), 1});
+    }
+
+    while (parser.Good()) {
+      if (parser.EatEqICase("aggregate")) {
+        std::string aggregate_value = GET_OR_RET(parser.TakeStr());
+        if (util::ToLower(aggregate_value) == "sum") {
+          aggregate_method_ = kAggregateSum;
+        } else if (util::ToLower(aggregate_value) == "min") {
+          aggregate_method_ = kAggregateMin;
+        } else if (util::ToLower(aggregate_value) == "max") {
+          aggregate_method_ = kAggregateMax;
+        } else {
+          return {Status::RedisParseErr, "aggregate param error"};
+        }
+      } else if (parser.EatEqICase("weights")) {
+        size_t k = 0;
+        while (k < numkeys_) {
+          auto weight = parser.TakeFloat();
+          if (!weight) {
+            return {Status::RedisParseErr, errValueIsNotFloat};
+          }
+          keys_weights_[k].weight = *weight;
+          k++;
+        }
+      } else if (parser.EatEqICase("withscores")) {
+        with_scores_ = true;
+      }
+    }
+    return Commander::Parse(args);
+  }
+
+  Status Execute(Server *svr, Connection *conn, std::string *output) override {
+    redis::ZSet zset_db(svr->storage, conn->GetNamespace());
+    std::vector<MemberScore> member_scores;
+    auto s = zset_db.Union(keys_weights_, aggregate_method_, nullptr, &member_scores);
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+    auto compare_score = [](const MemberScore &score1, const MemberScore &score2) {
+      if (score1.score == score2.score) {
+        return score1.member < score2.member;
+      }
+      return score1.score < score2.score;
+    };
+    std::sort(member_scores.begin(), member_scores.end(), compare_score);
+    output->append(redis::MultiLen(member_scores.size() * (with_scores_ ? 2 : 1)));
+    for (const auto &ms : member_scores) {
+      output->append(redis::BulkString(ms.member));
+      if (with_scores_) output->append(redis::BulkString(util::Float2String(ms.score)));
+    }
+    return Status::OK();
+  }
+
+  static CommandKeyRange Range(const std::vector<std::string> &args) {
+    int num_key = *ParseInt<int>(args[1], 10);
+    return {2, 1 + num_key, 1};

Review Comment:
   I'm not sure, but since `ZUnion` and `ZUnionStore` differs that `ZunionStore` has the `destination` argument, should their `Range` functions return different values (at least in one position)? @git-hulk @PragmaTwice 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@kvrocks.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org