Add support of the new command ZUNION

src/commands/cmd_zset.cc

@@ -1148,6 +1148,81 @@ class CommandZMScore : public Commander {
   }
 };
 
+class CommandZUnion : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    auto parse_result = ParseInt<int>(args[1], 10);
+    if (!parse_result) {
+      return {Status::RedisParseErr, errValueNotInteger};
+    }
+
+    numkeys_ = *parse_result;
+    if (numkeys_ > args.size() - 2) {
+      return {Status::RedisParseErr, errInvalidSyntax};
+    }
+
+    size_t key_iterator = 0;
+    while (key_iterator < numkeys_) {
+      keys_weights_.emplace_back(KeyWeight{args[key_iterator + 2], 1});
+      key_iterator++;
+    }
+
+    size_t option_iterator = 2 + numkeys_;
+    while (option_iterator < args.size()) {
+      if (util::ToLower(args[option_iterator]) == "aggregate" && option_iterator + 1 < args.size()) {
+        if (util::ToLower(args[option_iterator + 1]) == "sum") {
+          aggregate_method_ = kAggregateSum;
+        } else if (util::ToLower(args[option_iterator + 1]) == "min") {
+          aggregate_method_ = kAggregateMin;
+        } else if (util::ToLower(args[option_iterator + 1]) == "max") {
+          aggregate_method_ = kAggregateMax;
+        } else {
+          return {Status::RedisParseErr, "aggregate param error"};
+        }
+        option_iterator += 2;
+      } else if (util::ToLower(args[option_iterator]) == "weights" && option_iterator + numkeys_ < args.size()) {
+        size_t k = 0;
+        while (k < numkeys_) {
+          auto weight = ParseFloat(args[option_iterator + k + 1]);
+          if (!weight || std::isnan(*weight)) {
+            return {Status::RedisParseErr, "weight is not a double or out of range"};
+          }
+          keys_weights_[k].weight = *weight;
+          k++;
+        }
+        option_iterator += numkeys_ + 1;
+      } else if (util::ToLower(args[option_iterator]) == "withscores") {
+        with_scores_ = true;
+        option_iterator++;
+      } else {
+        return {Status::RedisParseErr, errInvalidSyntax};
+      }
+    }
+    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()};
+    }
+    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();
+  }
+
+ protected:
+  size_t numkeys_ = 0;
+  bool with_scores_ = false;
+  std::vector<KeyWeight> keys_weights_;
+  AggregateMethod aggregate_method_ = kAggregateSum;
+};
+
 class CommandZUnionStore : public Commander {
  public:
   Status Parse(const std::vector<std::string> &args) override {
@@ -1285,6 +1360,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandZAdd>("zadd", -4, "write", 1, 1, 1),
                         MakeCmdAttr<CommandZScore>("zscore", 3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandZMScore>("zmscore", -3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandZScan>("zscan", -3, "read-only", 1, 1, 1),
-                        MakeCmdAttr<CommandZUnionStore>("zunionstore", -4, "write", 1, 1, 1), )
+                        MakeCmdAttr<CommandZUnionStore>("zunionstore", -4, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandZUnion>("zunion", -4, "read-only", 1, 1, 1), )
 
 }  // namespace redis

src/types/redis_zset.cc

@@ -709,6 +709,16 @@ rocksdb::Status ZSet::InterStore(const Slice &dst, const std::vector<KeyWeight>
 
 rocksdb::Status ZSet::UnionStore(const Slice &dst, const std::vector<KeyWeight> &keys_weights,
                                  AggregateMethod aggregate_method, uint64_t *saved_cnt) {
+  *saved_cnt = 0;
+  std::vector<MemberScore> members;
+  auto s = Union(keys_weights, aggregate_method, saved_cnt, &members);
+  if (!s.ok()) return s;
+  *saved_cnt = members.size();
+  return Overwrite(dst, members);
+}
+
+rocksdb::Status ZSet::Union(const std::vector<KeyWeight> &keys_weights, AggregateMethod aggregate_method,
+                            uint64_t *saved_cnt, std::vector<MemberScore> *members) {
   if (saved_cnt) *saved_cnt = 0;
 
   std::map<std::string, double> dst_zset;
@@ -744,16 +754,15 @@ rocksdb::Status ZSet::UnionStore(const Slice &dst, const std::vector<KeyWeight>
       }
     }
   }
-  if (!dst_zset.empty()) {
-    std::vector<MemberScore> mscores;
-    mscores.reserve(dst_zset.size());
+  auto compare_score = [](const MemberScore &score1, const MemberScore &score2) { return score1.score < score2.score; };
+  if (members && !dst_zset.empty()) {
+    members->reserve(dst_zset.size());
     for (const auto &iter : dst_zset) {
-      mscores.emplace_back(MemberScore{iter.first, iter.second});
+      members->emplace_back(MemberScore{iter.first, iter.second});
     }
-    if (saved_cnt) *saved_cnt = mscores.size();
-    Overwrite(dst, mscores);
+    if (saved_cnt) *saved_cnt = members->size();
+    std::sort(members->begin(), members->end(), compare_score);
   }
-
   return rocksdb::Status::OK();
 }
 

src/types/redis_zset.h

@@ -110,6 +110,8 @@ class ZSet : public SubKeyScanner {
                              AggregateMethod aggregate_method, uint64_t *saved_cnt);
   rocksdb::Status UnionStore(const Slice &dst, const std::vector<KeyWeight> &keys_weights,
                              AggregateMethod aggregate_method, uint64_t *saved_cnt);
+  rocksdb::Status Union(const std::vector<KeyWeight> &keys_weights, AggregateMethod aggregate_method,
+                        uint64_t *saved_cnt, std::vector<MemberScore> *members);
   rocksdb::Status MGet(const Slice &user_key, const std::vector<Slice> &members, std::map<std::string, double> *scores);
   rocksdb::Status GetMetadata(const Slice &ns_key, ZSetMetadata *metadata);
 

tests/gocase/unit/type/zset/zset_test.go

@@ -1045,6 +1045,86 @@ func basicTests(t *testing.T, rdb *redis.Client, ctx context.Context, encoding s
 		require.Equal(t, []redis.Z{{1, "a"}, {2, "b"}, {3, "c"}, {3, "d"}}, rdb.ZRangeWithScores(ctx, "zsetc", 0, -1).Val())
 	})
 
+	t.Run(fmt.Sprintf("ZUNION basics - %s", encoding), func(t *testing.T) {
+		createZset(rdb, ctx, "zseta", []redis.Z{
+			{Score: 1, Member: "a"},
+			{Score: 2, Member: "b"},
+			{Score: 3, Member: "c"},
+		})
+		createZset(rdb, ctx, "zsetb", []redis.Z{
+			{Score: 1, Member: "b"},
+			{Score: 2, Member: "c"},
+			{Score: 3, Member: "d"},
+		})
+
+		zsetInt := []redis.Z{
+			{1, "a"},
+			{3, "b"},
+			{3, "d"},
+			{5, "c"}}
+		require.Equal(t, zsetInt, rdb.ZUnionWithScores(ctx, redis.ZStore{Keys: []string{"zseta", "zsetb"}}).Val())
+	})
+
+	t.Run(fmt.Sprintf("ZUNION with weights - %s", encoding), func(t *testing.T) {
+		createZset(rdb, ctx, "zseta", []redis.Z{
+			{Score: 1, Member: "a"},
+			{Score: 2, Member: "b"},
+			{Score: 3, Member: "c"},
+		})
+		createZset(rdb, ctx, "zsetb", []redis.Z{
+			{Score: 1, Member: "b"},
+			{Score: 2, Member: "c"},
+			{Score: 3, Member: "d"},
+		})
+
+		zsetInt := []redis.Z{
+			{2, "a"},
+			{7, "b"},
+			{9, "d"},
+			{12, "c"}}
+		require.Equal(t, zsetInt, rdb.ZUnionWithScores(ctx, redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}}).Val())
+	})
+
+	t.Run(fmt.Sprintf("ZUNION with AGGREGATE MIN - %s", encoding), func(t *testing.T) {
+		createZset(rdb, ctx, "zseta", []redis.Z{
+			{Score: 1, Member: "a"},
+			{Score: 2, Member: "b"},
+			{Score: 3, Member: "c"},
+		})
+		createZset(rdb, ctx, "zsetb", []redis.Z{
+			{Score: 1, Member: "b"},
+			{Score: 2, Member: "c"},
+			{Score: 3, Member: "d"},
+		})
+
+		zsetInt := []redis.Z{
+			{1, "a"},
+			{1, "b"},
+			{2, "c"},
+			{3, "d"}}
+		require.Equal(t, zsetInt, rdb.ZUnionWithScores(ctx, redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "min"}).Val())
+	})
+
+	t.Run(fmt.Sprintf("ZUNIONSTORE with AGGREGATE MAX - %s", encoding), func(t *testing.T) {
+		createZset(rdb, ctx, "zseta", []redis.Z{
+			{Score: 1, Member: "a"},
+			{Score: 2, Member: "b"},
+			{Score: 3, Member: "c"},
+		})
+		createZset(rdb, ctx, "zsetb", []redis.Z{
+			{Score: 1, Member: "b"},
+			{Score: 2, Member: "c"},
+			{Score: 3, Member: "d"},
+		})
+
+		zsetInt := []redis.Z{
+			{1, "a"},
+			{2, "b"},
+			{3, "c"},
+			{3, "d"}}
+		require.Equal(t, zsetInt, rdb.ZUnionWithScores(ctx, redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "max"}).Val())
+	})
+
 	t.Run(fmt.Sprintf("ZINTERSTORE basics - %s", encoding), func(t *testing.T) {
 		require.Equal(t, int64(2), rdb.ZInterStore(ctx, "zsetc", &redis.ZStore{Keys: []string{"zseta", "zsetb"}}).Val())
 		require.Equal(t, []redis.Z{{3, "b"}, {5, "c"}}, rdb.ZRangeWithScores(ctx, "zsetc", 0, -1).Val())