From 486e39e86eb2fd2f80ede87bd55d8e0c1b2c2a95 Mon Sep 17 00:00:00 2001
From: Tyson Andre <tyson.andre@uwaterloo.ca>
Date: Tue, 4 Aug 2020 10:49:33 -0400
Subject: [PATCH] Add a ZMSCORE command returning an array of scores. (#7593)

Syntax: `ZMSCORE KEY MEMBER [MEMBER ...]`

This is an extension of #2359
amended by Tyson Andre to work with the changed unstable API,
add more tests, and consistently return an array.

- It seemed as if it would be more likely to get reviewed
  after updating the implementation.

Currently, multi commands or lua scripting to call zscore multiple times
would almost definitely be less efficient than a native ZMSCORE
for the following reasons:

- Need to fetch the set from the string every time instead of reusing the C
  pointer.
- Using pipelining or multi-commands would result in more bytes sent by
  the client for the repeated `ZMSCORE KEY` sections.
- Need to specially encode the data and decode it from the client
  for lua-based solutions.
- The fastest solution I've seen for large sets(thousands or millions)
  involves lua and a variadic ZADD, then a ZINTERSECT, then a ZRANGE 0 -1,
  then UNLINK of a temporary set (or lua). This is still inefficient.

Co-authored-by: Tyson Andre <tysonandre775@hotmail.com>
---
 src/server.c              |  4 +++
 src/server.h              |  1 +
 src/t_zset.c              | 18 +++++++++++++
 tests/support/cluster.tcl |  2 +-
 tests/unit/type/zset.tcl  | 53 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/src/server.c b/src/server.c
index afe5094c5..9d915ac3b 100644
--- a/src/server.c
+++ b/src/server.c
@@ -470,6 +470,10 @@ struct redisCommand redisCommandTable[] = {
      "read-only fast @sortedset",
      0,NULL,1,1,1,0,0,0},
 
+    {"zmscore",zmscoreCommand,-3,
+     "read-only fast @sortedset",
+     0,NULL,1,1,1,0,0,0},
+
     {"zrank",zrankCommand,3,
      "read-only fast @sortedset",
      0,NULL,1,1,1,0,0,0},
diff --git a/src/server.h b/src/server.h
index adfac3293..018dd4607 100644
--- a/src/server.h
+++ b/src/server.h
@@ -2321,6 +2321,7 @@ void zrevrangeCommand(client *c);
 void zcardCommand(client *c);
 void zremCommand(client *c);
 void zscoreCommand(client *c);
+void zmscoreCommand(client *c);
 void zremrangebyscoreCommand(client *c);
 void zremrangebylexCommand(client *c);
 void zpopminCommand(client *c);
diff --git a/src/t_zset.c b/src/t_zset.c
index 16131ee85..87a5a882e 100644
--- a/src/t_zset.c
+++ b/src/t_zset.c
@@ -3082,6 +3082,24 @@ void zscoreCommand(client *c) {
     }
 }
 
+void zmscoreCommand(client *c) {
+    robj *key = c->argv[1];
+    robj *zobj;
+    double score;
+    zobj = lookupKeyRead(c->db,key);
+    if (zobj != NULL && checkType(c,zobj,OBJ_ZSET)) return;
+
+    addReplyArrayLen(c,c->argc - 2);
+    for (int j = 2; j < c->argc; j++) {
+        /* Treat a missing set the same way as an empty set */
+        if (zobj == NULL || zsetScore(zobj,c->argv[j]->ptr,&score) == C_ERR) {
+            addReplyNull(c);
+        } else {
+            addReplyDouble(c,score);
+        }
+    }
+}
+
 void zrankGenericCommand(client *c, int reverse) {
     robj *key = c->argv[1];
     robj *ele = c->argv[2];
diff --git a/tests/support/cluster.tcl b/tests/support/cluster.tcl
index 64b079ff8..dced8d864 100644
--- a/tests/support/cluster.tcl
+++ b/tests/support/cluster.tcl
@@ -27,7 +27,7 @@ set ::redis_cluster::plain_commands {
     sadd srem sismember scard spop srandmember smembers sscan zadd
     zincrby zrem zremrangebyscore zremrangebyrank zremrangebylex zrange
     zrangebyscore zrevrangebyscore zrangebylex zrevrangebylex zcount
-    zlexcount zrevrange zcard zscore zrank zrevrank zscan hset hsetnx
+    zlexcount zrevrange zcard zscore zmscore zrank zrevrank zscan hset hsetnx
     hget hmset hmget hincrby hincrbyfloat hdel hlen hkeys hvals
     hgetall hexists hscan incrby decrby incrbyfloat getset move
     expire expireat pexpire pexpireat type ttl pttl persist restore
diff --git a/tests/unit/type/zset.tcl b/tests/unit/type/zset.tcl
index a8c817f6e..d098840b0 100644
--- a/tests/unit/type/zset.tcl
+++ b/tests/unit/type/zset.tcl
@@ -780,6 +780,44 @@ start_server {tags {"zset"}} {
             set oldscore $score
         }
     }
+    
+    test {ZMSCORE retrieve} {
+        r del zmscoretest
+        r zadd zmscoretest 10 x
+        r zadd zmscoretest 20 y
+        
+        r zmscore zmscoretest x y
+    } {10 20}
+
+    test {ZMSCORE retrieve from empty set} {
+        r del zmscoretest
+        
+        r zmscore zmscoretest x y
+    } {{} {}}
+    
+    test {ZMSCORE retrieve with missing member} {
+        r del zmscoretest
+        r zadd zmscoretest 10 x
+        
+        r zmscore zmscoretest x y
+    } {10 {}}
+
+    test {ZMSCORE retrieve single member} {
+        r del zmscoretest
+        r zadd zmscoretest 10 x
+        r zadd zmscoretest 20 y
+        
+        r zmscore zmscoretest x
+    } {10}
+
+    test {ZMSCORE retrieve requires one or more members} {
+        r del zmscoretest
+        r zadd zmscoretest 10 x
+        r zadd zmscoretest 20 y
+        
+        catch {r zmscore zmscoretest} e
+        assert_match {*ERR*wrong*number*arg*} $e
+    }
 
     test "ZSET commands don't accept the empty strings as valid score" {
         assert_error "*not*float*" {r zadd myzset "" abc}
@@ -815,6 +853,21 @@ start_server {tags {"zset"}} {
             }
         }
 
+        test "ZMSCORE - $encoding" {
+            r del zscoretest
+            set aux {}
+            for {set i 0} {$i < $elements} {incr i} {
+                set score [expr rand()]
+                lappend aux $score
+                r zadd zscoretest $score $i
+            }
+
+            assert_encoding $encoding zscoretest
+            for {set i 0} {$i < $elements} {incr i} {
+                assert_equal [lindex $aux $i] [r zmscore zscoretest $i]
+            }
+        }
+
         test "ZSCORE after a DEBUG RELOAD - $encoding" {
             r del zscoretest
             set aux {}