From 29fb499d9b09e6b626805107f3a187bd367daf31 Mon Sep 17 00:00:00 2001 From: Logan Raarup Date: Fri, 7 Oct 2022 13:21:06 +0200 Subject: [PATCH] feat(server): Implement ZMSCORE command #357 (#367) * feat(server): Implement ZMSCORE command (#357) Signed-off-by: Logan Raarup --- CONTRIBUTORS.md | 1 + src/server/zset_family.cc | 56 ++++++++++++++++++++++++++++++++++ src/server/zset_family.h | 4 +++ src/server/zset_family_test.cc | 9 ++++++ 4 files changed, 70 insertions(+) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 9a0b0fb..2abda5f 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -4,6 +4,7 @@ * **[Philipp Born](https://github.com/tamcore)** * Helm Chart * **[Braydn Moore](https://github.com/braydnm)** +* **[Logan Raarup](https://github.com/logandk)** * **[Ryan Russell](https://github.com/ryanrussell)** * Docs & Code Readability * **[Ali-Akber Saifee](https://github.com/alisaifee)** diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc index dd5e021..2d66bd7 100644 --- a/src/server/zset_family.cc +++ b/src/server/zset_family.cc @@ -1274,6 +1274,35 @@ void ZSetFamily::ZScore(CmdArgList args, ConnectionContext* cntx) { } } +void ZSetFamily::ZMScore(CmdArgList args, ConnectionContext* cntx) { + string_view key = ArgS(args, 1); + + absl::InlinedVector members(args.size() - 2); + for (size_t i = 2; i < args.size(); ++i) { + members[i - 2] = ArgS(args, i); + } + + auto cb = [&](Transaction* t, EngineShard* shard) { + return OpMScore(t->GetOpArgs(shard), key, members); + }; + + OpResult result = cntx->transaction->ScheduleSingleHopT(std::move(cb)); + + if (result.status() == OpStatus::WRONG_TYPE) { + return (*cntx)->SendError(kWrongTypeErr); + } + + (*cntx)->StartArray(result->size()); + const MScoreResponse& array = result.value(); + for (const auto& p : array) { + if (p) { + (*cntx)->SendDouble(*p); + } else { + (*cntx)->SendNull(); + } + } +} + void ZSetFamily::ZScan(CmdArgList args, ConnectionContext* cntx) { string_view key = ArgS(args, 1); string_view token = ArgS(args, 2); @@ -1603,6 +1632,32 @@ OpResult ZSetFamily::OpScore(const OpArgs& op_args, string_view key, str return score; } +OpResult ZSetFamily::OpMScore(const OpArgs& op_args, string_view key, ArgSlice members) { + OpResult res_it = op_args.shard->db_slice().Find(op_args.db_cntx, key, OBJ_ZSET); + if (!res_it) + return res_it.status(); + + MScoreResponse scores(members.size()); + + robj* zobj = res_it.value()->second.AsRObj(); + sds& tmp_str = op_args.shard->tmp_str1; + + for (size_t i = 0; i < members.size(); i++) { + const auto& m = members[i]; + + tmp_str = sdscpylen(tmp_str, m.data(), m.size()); + double score; + int retval = zsetScore(zobj, tmp_str, &score); + if (retval == C_OK) { + scores[i] = score; + } else { + scores[i] = std::nullopt; + } + } + + return scores; +} + auto ZSetFamily::OpRange(const ZRangeSpec& range_spec, const OpArgs& op_args, string_view key) -> OpResult { OpResult res_it = op_args.shard->db_slice().Find(op_args.db_cntx, key, OBJ_ZSET); @@ -1808,6 +1863,7 @@ void ZSetFamily::Register(CommandRegistry* registry) { << CI{"ZRANGEBYLEX", CO::READONLY, -4, 1, 1, 1}.HFUNC(ZRangeByLex) << CI{"ZRANGEBYSCORE", CO::READONLY, -4, 1, 1, 1}.HFUNC(ZRangeByScore) << CI{"ZSCORE", CO::READONLY | CO::FAST, 3, 1, 1, 1}.HFUNC(ZScore) + << CI{"ZMSCORE", CO::READONLY | CO::FAST, -3, 1, 1, 1}.HFUNC(ZMScore) << CI{"ZREMRANGEBYRANK", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByRank) << CI{"ZREMRANGEBYSCORE", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByScore) << CI{"ZREMRANGEBYLEX", CO::WRITE, 4, 1, 1, 1}.HFUNC(ZRemRangeByLex) diff --git a/src/server/zset_family.h b/src/server/zset_family.h index ced187e..9ee8577 100644 --- a/src/server/zset_family.h +++ b/src/server/zset_family.h @@ -62,6 +62,7 @@ class ZSetFamily { static void ZRank(CmdArgList args, ConnectionContext* cntx); static void ZRem(CmdArgList args, ConnectionContext* cntx); static void ZScore(CmdArgList args, ConnectionContext* cntx); + static void ZMScore(CmdArgList args, ConnectionContext* cntx); static void ZRangeByLex(CmdArgList args, ConnectionContext* cntx); static void ZRangeByScore(CmdArgList args, ConnectionContext* cntx); static void ZRemRangeByRank(CmdArgList args, ConnectionContext* cntx); @@ -89,6 +90,9 @@ class ZSetFamily { static OpResult OpRem(const OpArgs& op_args, std::string_view key, ArgSlice members); static OpResult OpScore(const OpArgs& op_args, std::string_view key, std::string_view member); + using MScoreResponse = std::vector>; + static OpResult OpMScore(const OpArgs& op_args, std::string_view key, + ArgSlice members); static OpResult OpRange(const ZRangeSpec& range_spec, const OpArgs& op_args, std::string_view key); static OpResult OpRemRange(const OpArgs& op_args, std::string_view key, diff --git a/src/server/zset_family_test.cc b/src/server/zset_family_test.cc index 4d31d0c..141394c 100644 --- a/src/server/zset_family_test.cc +++ b/src/server/zset_family_test.cc @@ -64,6 +64,15 @@ TEST_F(ZSetFamilyTest, ZRem) { EXPECT_THAT(Run({"zrange", "x", "(-inf", "(+inf", "byscore"}), "a"); } +TEST_F(ZSetFamilyTest, ZMScore) { + Run({"zadd", "zms", "3.14", "a"}); + Run({"zadd", "zms", "42", "another"}); + + auto resp = Run({"zmscore", "zms", "another", "a", "nofield"}); + ASSERT_EQ(RespExpr::ARRAY, resp.type); + EXPECT_THAT(resp.GetVec(), ElementsAre("42", "3.14", ArgType(RespExpr::NIL))); +} + TEST_F(ZSetFamilyTest, ZRangeRank) { Run({"zadd", "x", "1.1", "a", "2.1", "b"}); EXPECT_THAT(Run({"zrangebyscore", "x", "0", "(1.1"}), ArrLen(0));