Add support of the new command ZUNION

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};
+  }
+
+ 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 {
@@ -1211,6 +1283,11 @@ class CommandZUnionStore : public Commander {
     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};
+  }
+
  protected:
   size_t numkeys_ = 0;
   std::vector<KeyWeight> keys_weights_;
@@ -1285,6 +1362,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", CommandZUnionStore::Range),
+                        MakeCmdAttr<CommandZUnion>("zunion", -4, "read-only", CommandZUnion::Range), )
 
 }  // 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,13 @@ rocksdb::Status ZSet::UnionStore(const Slice &dst, const std::vector<KeyWeight>
       }
     }
   }
-  if (!dst_zset.empty()) {
-    std::vector<MemberScore> mscores;
-    mscores.reserve(dst_zset.size());
+  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();
   }
-
   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/cppunit/command_parser_test.cc

@@ -26,9 +26,10 @@ TEST(CommandParser, Parse) {
   // [ HELLO i1 v1 | HI v2 ] [X i2 | Y]
   std::vector<std::string> c1{"hello", "1", "a"}, c2{"hi", "b"}, c3{"hi", "c", "x", "2"}, c4{"hello", "3", "d", "y"},
       c5{"hi", "e", "y"}, c6{"y"}, d1{"hello"}, d2{"hi"}, d3{"hello", "no-int"}, d4{"x", "1", "y"},
-      d5{"hello", "1", "v", "hi", "v"}, d6{"hello", "1"};
+      d5{"hello", "1", "v", "hi", "v"}, d6{"hello", "1"}, d7{"float_test", "1.1"};
 
   std::int64_t i1 = 0, i2 = 0;
+  double f1 = 0;
   std::string v1, v2;
   std::string_view hflag, xflag;
 
@@ -44,6 +45,8 @@ TEST(CommandParser, Parse) {
         i2 = GET_OR_RET(parser.TakeInt());
       } else if (parser.EatEqICaseFlag("y", xflag)) {
         // pass
+      } else if (parser.EatEqICaseFlag("float_test", hflag)) {
+        f1 = GET_OR_RET(parser.TakeFloat());
       } else {
         return parser.InvalidSyntax();
       }
@@ -76,6 +79,8 @@ TEST(CommandParser, Parse) {
   ASSERT_FALSE(parse(CommandParser(d4)));
   ASSERT_FALSE(parse(CommandParser(d5)));
   ASSERT_FALSE(parse(CommandParser(d6)));
+  ASSERT_TRUE(parse(CommandParser(d7)));
+  ASSERT_EQ(f1, 1.1);
 }
 
 TEST(CommandParser, ParseMove) {

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())