Add support of the command ZDIFF and ZDIFFSTORE

src/commands/cmd_zset.cc

@@ -1357,6 +1357,107 @@ class CommandZScan : public CommandSubkeyScanBase {
   }
 };
 
+class CommandZDiff : 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 j = 0;
+    while (j < numkeys_) {
+      keys_.emplace_back(args[j + 2]);
+      j++;
+    }
+
+    if (auto i = 2 + numkeys_; i < args.size()) {
+      if (util::ToLower(args[i]) == "withscores") {
+        with_scores_ = true;
+      }
+    }
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::ZSet zset_db(srv->storage, conn->GetNamespace());
+
+    std::vector<MemberScore> members_with_scores;
+    auto s = zset_db.Diff(keys_, &members_with_scores);
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    output->append(redis::MultiLen(members_with_scores.size() * (with_scores_ ? 2 : 1)));
+    for (const auto &ms : members_with_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, 2 + num_key, 1};
+  }
+
+ protected:
+  size_t numkeys_ {0};
+  std::vector<rocksdb::Slice> keys_;
+  bool with_scores_ {false};
+};
+
+class CommandZDiffStore : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    auto parse_result = ParseInt<int>(args[2], 10);
+    if (!parse_result) return {Status::RedisParseErr, errValueNotInteger};
+
+    numkeys_ = *parse_result;
+    if (numkeys_ > args.size() - 3) return {Status::RedisParseErr, errInvalidSyntax};
+
+    size_t j = 0;
+    while (j < numkeys_) {
+      keys_.emplace_back(args[j + 3]);
+      j++;
+    }
+
+    if (auto i = 2 + numkeys_; i < args.size()) {
+      if (util::ToLower(args[i]) == "withscores") {
+        with_scores_ = true;
+      }
+    }
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::ZSet zset_db(srv->storage, conn->GetNamespace());
+
+    uint64_t stored_count;
+    LOG(INFO) << args_[1];
+    auto s = zset_db.DiffStore(args_[1], keys_, &stored_count);
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+    *output = redis::Integer(stored_count);
+    return Status::OK();
+  }
+
+  static CommandKeyRange Range(const std::vector<std::string> &args) {
+    int num_key = *ParseInt<int>(args[1], 10);
+    return {3, 2 + num_key, 1};
+  }
+
+ protected:
+  size_t numkeys_ {0};
+  std::vector<rocksdb::Slice> keys_;
+  bool with_scores_ {false};
+};
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandZAdd>("zadd", -4, "write", 1, 1, 1),
                         MakeCmdAttr<CommandZCard>("zcard", 2, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandZCount>("zcount", 4, "read-only", 1, 1, 1),
@@ -1388,6 +1489,7 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandZAdd>("zadd", -4, "write", 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", CommandZUnionStore::Range),
-                        MakeCmdAttr<CommandZUnion>("zunion", -3, "read-only", CommandZUnion::Range), )
-
+                        MakeCmdAttr<CommandZUnion>("zunion", -3, "read-only", CommandZUnion::Range),                        
+                        MakeCmdAttr<CommandZDiff>("zdiff", -3, "read-only", CommandZDiff::Range),
+                        MakeCmdAttr<CommandZDiffStore>("zdiffstore", -3, "read-only", CommandZDiffStore::Range), )
 }  // namespace redis

src/types/redis_zset.cc

@@ -851,4 +851,46 @@ rocksdb::Status ZSet::MGet(const Slice &user_key, const std::vector<Slice> &memb
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status ZSet::Diff(const std::vector<Slice> &keys, MemberScores *members) {
+  std::vector<std::string> lock_keys;
+  lock_keys.reserve(keys.size());
+  for (const auto key : keys) {
+    std::string ns_key = AppendNamespacePrefix(key);
+    lock_keys.emplace_back(std::move(ns_key));
+  }
+  MultiLockGuard guard(storage_->GetLockManager(), lock_keys);
+
+  members->clear();
+  MemberScores source_member_scores;
+  RangeScoreSpec spec;
+  uint64_t size {0};
+  auto s = RangeByScore(keys[0], spec, &source_member_scores, &size);
+  if (!s.ok()) return s;
+
+  std::map<std::string, bool> exclude_members {};
+  MemberScores target_member_scores {};
+  for (size_t i = 1; i < keys.size(); i++) {
+    uint64_t size {0};
+    s = RangeByScore(keys[i], spec, &target_member_scores, &size);
+    if (!s.ok()) return s;
+    for (const auto &member_score : target_member_scores) {
+      exclude_members[member_score.member] = true;
+    }
+  }
+  for (const auto &member_score : source_member_scores) {
+    if (exclude_members.find(member_score.member) == exclude_members.end()) {
+      members->push_back(member_score);
+    }
+  }
+  return rocksdb::Status::OK();
+}
+
+rocksdb::Status ZSet::DiffStore(const Slice &dst, const std::vector<Slice> &keys, uint64_t *stored_count) {
+  MemberScores mscores;
+  auto s = Diff(keys, &mscores);
+  if (!s.ok()) return s;
+  *stored_count = mscores.size();
+  return Overwrite(dst, mscores);
+}
+
 }  // namespace redis

src/types/redis_zset.h

@@ -116,6 +116,8 @@ class ZSet : public SubKeyScanner {
                              AggregateMethod aggregate_method, uint64_t *saved_cnt);
   rocksdb::Status Union(const std::vector<KeyWeight> &keys_weights, AggregateMethod aggregate_method,
                         std::vector<MemberScore> *members);
+  rocksdb::Status Diff(const std::vector<Slice> &keys, MemberScores *members);
+  rocksdb::Status DiffStore(const Slice &dst, const std::vector<Slice> &keys, uint64_t *stored_count);  
   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/types/zset_test.cc

@@ -433,3 +433,54 @@ TEST_F(RedisZSetTest, Rank) {
   }
   auto s = zset_->Del(key_);
 }
+
+TEST_F(RedisZSetTest, Diff) {
+  uint64_t ret = 0;
+
+  std::string k1 = "key1";
+  std::vector<rocksdb::Slice> k1_fields_ = {"a", "b", "c", "d"};
+  std::vector<double> k1_scores_ = {-100.1, -100.1, 0, 1.234};
+  std::vector<MemberScore> k1_mscores;
+  for (size_t i = 0; i < k1_fields_.size(); i++) {
+    k1_mscores.emplace_back(MemberScore{k1_fields_[i].ToString(), k1_scores_[i]});
+  }
+
+  std::string k2 = "key2";
+  std::vector<rocksdb::Slice> k2_fields_ = {"c"};
+  std::vector<double> k2_scores_ = {-150.1};
+  std::vector<MemberScore> k2_mscores;
+  for (size_t i = 0; i < k2_fields_.size(); i++) {
+    k2_mscores.emplace_back(MemberScore{k2_fields_[i].ToString(), k2_scores_[i]});
+  }
+
+  std::string k3 = "key3";
+  std::vector<rocksdb::Slice> k3_fields_ = {"a", "c", "e"};
+  std::vector<double> k3_scores_ = {-1000.1, -100.1, 8000.9};
+  std::vector<MemberScore> k3_mscores;
+  for (size_t i = 0; i < k3_fields_.size(); i++) {
+    k3_mscores.emplace_back(MemberScore{k3_fields_[i].ToString(), k3_scores_[i]});
+  }
+
+  auto s = zset_->Add(k1, ZAddFlags::Default(), &k1_mscores, &ret);
+  EXPECT_EQ(ret, 4);
+  zset_->Add(k2, ZAddFlags::Default(), &k2_mscores, &ret);
+  EXPECT_EQ(ret, 1);
+  zset_->Add(k3, ZAddFlags::Default(), &k3_mscores, &ret);
+  EXPECT_EQ(ret, 3);
+
+  std::vector<MemberScore> mscores;
+  zset_->Diff({k1, k2, k3}, &mscores);
+
+  EXPECT_EQ(2, mscores.size());
+  std::vector<MemberScore> expected_mscores = {{"b", -100.1}, {"d", 1.234}};
+  int index = 0;
+  for (auto mscore : expected_mscores) {
+    EXPECT_EQ(mscore.member, mscores[index].member);
+    EXPECT_EQ(mscore.score, mscores[index].score);
+    index++;
+  }
+
+  s = zset_->Del(k1);
+  s = zset_->Del(k2);
+  s = zset_->Del(k3);
+}

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

@@ -1330,6 +1330,79 @@ func basicTests(t *testing.T, rdb *redis.Client, ctx context.Context, encoding s
 				Weights: []float64{math.NaN(), math.NaN()}},
 			).Err(), ".*weight.*not.*double.*")
 		})
+
+		t.Run(fmt.Sprintf("ZDIFF with two sets - %s", encoding), func(t *testing.T) {
+			createZset(rdb, ctx, "zseta", []redis.Z{
+				{Score: 1, Member: "a"},
+				{Score: 2, Member: "b"},
+				{Score: 3, Member: "c"},
+				{Score: 3, Member: "d"},
+				{Score: 4, Member: "e"},
+			})
+			createZset(rdb, ctx, "zsetb", []redis.Z{
+				{Score: 1, Member: "b"},
+				{Score: 2, Member: "c"},
+				{Score: 4, Member: "f"},
+			})
+			cmd := rdb.ZDiff(ctx, "zseta", "zsetb")
+			require.NoError(t, cmd.Err())
+			sort.Strings(cmd.Val())
+			require.EqualValues(t, []string{"a", "d", "e"}, cmd.Val())
+		})
+
+		t.Run(fmt.Sprintf("ZDIFF with three sets - %s", encoding), func(t *testing.T) {
+			createZset(rdb, ctx, "zseta", []redis.Z{
+				{Score: 1, Member: "a"},
+				{Score: 2, Member: "b"},
+				{Score: 3, Member: "c"},
+				{Score: 3, Member: "d"},
+				{Score: 4, Member: "e"},
+			})
+			createZset(rdb, ctx, "zsetb", []redis.Z{
+				{Score: 1, Member: "b"},
+				{Score: 2, Member: "c"},
+				{Score: 4, Member: "f"},
+			})
+			createZset(rdb, ctx, "zsetc", []redis.Z{
+				{Score: 3, Member: "c"},
+				{Score: 3, Member: "d"},
+				{Score: 5, Member: "e"},
+			})
+			cmd := rdb.ZDiff(ctx, "zseta", "zsetb", "zsetc")
+			require.NoError(t, cmd.Err())
+			sort.Strings(cmd.Val())
+			require.EqualValues(t, []string{"a"}, cmd.Val())
+		})
+
+		t.Run(fmt.Sprintf("ZDIFF with three sets with scores - %s", encoding), func(t *testing.T) {
+			createZset(rdb, ctx, "zseta", []redis.Z{
+				{Score: 1, Member: "a"},
+				{Score: 2, Member: "b"},
+				{Score: 3, Member: "c"},
+				{Score: 3, Member: "d"},
+				{Score: 4, Member: "e"},
+			})
+			createZset(rdb, ctx, "zsetb", []redis.Z{
+				{Score: 1, Member: "b"},
+				{Score: 2, Member: "c"},
+				{Score: 4, Member: "f"},
+			})
+			createZset(rdb, ctx, "zsetc", []redis.Z{
+				{Score: 4, Member: "c"},
+				{Score: 5, Member: "e"},
+			})
+			cmd := rdb.ZDiffWithScores(ctx, "zseta", "zsetb", "zsetc")
+			require.NoError(t, cmd.Err())
+			require.EqualValues(t, []redis.Z([]redis.Z{{Score: 1, Member: "a"}, {Score: 3, Member: "d"}}), cmd.Val())
+		})
+
+		// t.Run("ZDIFFSTORE with three sets - ", func(t *testing.T) {
+		// 	require.NoError(t, rdb.ZDiffStore(ctx, "setres", "set1", "set4", "set5").Err())
+		// 	cmd := rdb.SMembers(ctx, "setres")
+		// 	require.NoError(t, cmd.Err())
+		// 	sort.Strings(cmd.Val())
+		// 	require.EqualValues(t, []string{"1", "2", "3", "4"}, cmd.Val())
+		// })
 	}
 }