Support Z[REV]RANK command with score option (#3364)

* Support z[rev]rank with score command fixes #3350 #3349 * Update src/main/java/redis/clients/jedis/BuilderFactory.java Co-authored-by: M Sazzadul Hoque <[email protected]> --------- Co-authored-by: M Sazzadul Hoque <[email protected]>
redis · Apr 21, 2023 · b2c5acf · b2c5acf
1 parent 8b60f29
commit b2c5acf
Show file tree

Hide file tree

Showing 14 changed files with 255 additions and 1 deletion.
diff --git a/src/main/java/redis/clients/jedis/BuilderFactory.java b/src/main/java/redis/clients/jedis/BuilderFactory.java
@@ -389,6 +389,22 @@ public String toString() {
     }
   };
 
+  public static final Builder<KeyValue<Long, Double>> ZRANK_WITHSCORE_PAIR = new Builder<KeyValue<Long, Double>>() {
+    @Override
+    public KeyValue<Long, Double> build(Object data) {
+      if (data == null) {
+        return null;
+      }
+      List<Object> l = (List<Object>) data;
+      return new KeyValue<>(LONG.build(l.get(0)), DOUBLE.build(l.get(1)));
+    }
+
+    @Override
+    public String toString() {
+      return "KeyValue<Long, Double>";
+    }
+  };
+
   public static final Builder<KeyValue<String, List<String>>> KEYED_STRING_LIST
       = new Builder<KeyValue<String, List<String>>>() {
     @Override

diff --git a/src/main/java/redis/clients/jedis/CommandObjects.java b/src/main/java/redis/clients/jedis/CommandObjects.java
@@ -1380,6 +1380,14 @@ public final CommandObject<Long> zrevrank(String key, String member) {
     return new CommandObject<>(commandArguments(ZREVRANK).key(key).add(member), BuilderFactory.LONG);
   }
 
+  public final CommandObject<KeyValue<Long, Double>> zrankWithScore(String key, String member) {
+    return new CommandObject<>(commandArguments(ZRANK).key(key).add(member).add(WITHSCORE), BuilderFactory.ZRANK_WITHSCORE_PAIR);
+  }
+
+  public final CommandObject<KeyValue<Long, Double>> zrevrankWithScore(String key, String member) {
+    return new CommandObject<>(commandArguments(ZREVRANK).key(key).add(member).add(WITHSCORE), BuilderFactory.ZRANK_WITHSCORE_PAIR);
+  }
+
   public final CommandObject<Long> zrank(byte[] key, byte[] member) {
     return new CommandObject<>(commandArguments(ZRANK).key(key).add(member), BuilderFactory.LONG);
   }
@@ -1388,6 +1396,14 @@ public final CommandObject<Long> zrevrank(byte[] key, byte[] member) {
     return new CommandObject<>(commandArguments(ZREVRANK).key(key).add(member), BuilderFactory.LONG);
   }
 
+  public final CommandObject<KeyValue<Long, Double>> zrankWithScore(byte[] key, byte[] member) {
+    return new CommandObject<>(commandArguments(ZRANK).key(key).add(member).add(WITHSCORE), BuilderFactory.ZRANK_WITHSCORE_PAIR);
+  }
+
+  public final CommandObject<KeyValue<Long, Double>> zrevrankWithScore(byte[] key, byte[] member) {
+    return new CommandObject<>(commandArguments(ZREVRANK).key(key).add(member).add(WITHSCORE), BuilderFactory.ZRANK_WITHSCORE_PAIR);
+  }
+
   public final CommandObject<String> zrandmember(String key) {
     return new CommandObject<>(commandArguments(ZRANDMEMBER).key(key), BuilderFactory.STRING);
   }

diff --git a/src/main/java/redis/clients/jedis/Jedis.java b/src/main/java/redis/clients/jedis/Jedis.java
@@ -2130,6 +2130,32 @@ public Long zrevrank(final byte[] key, final byte[] member) {
     return connection.executeCommand(commandObjects.zrevrank(key, member));
   }
 
+  /**
+   * Returns the rank and the score of member in the sorted set stored at key, with the scores
+   * ordered from low to high.
+   * @param key the key
+   * @param member the member
+   * @return the KeyValue contains rank and score.
+   */
+  @Override
+  public KeyValue<Long, Double> zrankWithScore(byte[] key, byte[] member) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  /**
+   * Returns the rank and the score of member in the sorted set stored at key, with the scores
+   * ordered from high to low.
+   * @param key the key
+   * @param member the member
+   * @return the KeyValue contains rank and score.
+   */
+  @Override
+  public KeyValue<Long, Double> zrevrankWithScore(byte[] key, byte[] member) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public List<byte[]> zrevrange(final byte[] key, final long start, final long stop) {
     checkIsInMultiOrPipeline();
@@ -6559,6 +6585,32 @@ public Long zrevrank(final String key, final String member) {
     return connection.executeCommand(commandObjects.zrevrank(key, member));
   }
 
+  /**
+   * Returns the rank and the score of member in the sorted set stored at key, with the scores
+   * ordered from low to high.
+   * @param key the key
+   * @param member the member
+   * @return the KeyValue contains rank and score.
+   */
+  @Override
+  public KeyValue<Long, Double> zrankWithScore(String key, String member) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  /**
+   * Returns the rank and the score of member in the sorted set stored at key, with the scores
+   * ordered from high to low.
+   * @param key the key
+   * @param member the member
+   * @return the KeyValue contains rank and score.
+   */
+  @Override
+  public KeyValue<Long, Double> zrevrankWithScore(String key, String member) {
+    checkIsInMultiOrPipeline();
+    return connection.executeCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public List<String> zrevrange(final String key, final long start, final long stop) {
     checkIsInMultiOrPipeline();

diff --git a/src/main/java/redis/clients/jedis/MultiNodePipelineBase.java b/src/main/java/redis/clients/jedis/MultiNodePipelineBase.java
@@ -976,6 +976,16 @@ public Response<Long> zrevrank(String key, String member) {
     return appendCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public Response<KeyValue<Long, Double>> zrankWithScore(String key, String member) {
+    return appendCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public Response<KeyValue<Long, Double>> zrevrankWithScore(String key, String member) {
+    return appendCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public Response<List<String>> zrange(String key, long start, long stop) {
     return appendCommand(commandObjects.zrange(key, start, stop));
@@ -2724,6 +2734,16 @@ public Response<Long> zrevrank(byte[] key, byte[] member) {
     return appendCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public Response<KeyValue<Long, Double>> zrankWithScore(byte[] key, byte[] member) {
+    return appendCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public Response<KeyValue<Long, Double>> zrevrankWithScore(byte[] key, byte[] member) {
+    return appendCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public Response<List<byte[]>> zrange(byte[] key, long start, long stop) {
     return appendCommand(commandObjects.zrange(key, start, stop));

diff --git a/src/main/java/redis/clients/jedis/Pipeline.java b/src/main/java/redis/clients/jedis/Pipeline.java
@@ -928,6 +928,16 @@ public Response<Long> zrevrank(String key, String member) {
     return appendCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public Response<KeyValue<Long, Double>> zrankWithScore(String key, String member) {
+    return appendCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public Response<KeyValue<Long, Double>> zrevrankWithScore(String key, String member) {
+    return appendCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public Response<List<String>> zrange(String key, long start, long stop) {
     return appendCommand(commandObjects.zrange(key, start, stop));
@@ -2675,6 +2685,16 @@ public Response<Long> zrevrank(byte[] key, byte[] member) {
     return appendCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public Response<KeyValue<Long, Double>> zrankWithScore(byte[] key, byte[] member) {
+    return appendCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public Response<KeyValue<Long, Double>> zrevrankWithScore(byte[] key, byte[] member) {
+    return appendCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public Response<List<byte[]>> zrange(byte[] key, long start, long stop) {
     return appendCommand(commandObjects.zrange(key, start, stop));

diff --git a/src/main/java/redis/clients/jedis/Protocol.java b/src/main/java/redis/clients/jedis/Protocol.java
@@ -253,7 +253,7 @@ public byte[] getRaw() {
 
   public static enum Keyword implements Rawable {
 
-    AGGREGATE, ALPHA, BY, GET, LIMIT, NO, NOSORT, ONE, SET, STORE, WEIGHTS, WITHSCORES, RESETSTAT,
+    AGGREGATE, ALPHA, BY, GET, LIMIT, NO, NOSORT, ONE, SET, STORE, WEIGHTS, WITHSCORE, WITHSCORES, RESETSTAT,
     REWRITE, RESET, FLUSH, EXISTS, LOAD, LEN, HELP, SCHEDULE, MATCH, COUNT, TYPE, KEYS, REFCOUNT,
     ENCODING, IDLETIME, FREQ, REPLACE, GETNAME, SETNAME, SETINFO, LIST, ID, KILL, PAUSE, UNPAUSE, UNBLOCK,
     STREAMS, CREATE, MKSTREAM, SETID, DESTROY, DELCONSUMER, MAXLEN, GROUP, IDLE, TIME, BLOCK, NOACK,

diff --git a/src/main/java/redis/clients/jedis/TransactionBase.java b/src/main/java/redis/clients/jedis/TransactionBase.java
@@ -1025,6 +1025,16 @@ public Response<Long> zrevrank(String key, String member) {
     return appendCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public Response<KeyValue<Long, Double>> zrankWithScore(String key, String member) {
+    return appendCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public Response<KeyValue<Long, Double>> zrevrankWithScore(String key, String member) {
+    return appendCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public Response<List<String>> zrange(String key, long start, long stop) {
     return appendCommand(commandObjects.zrange(key, start, stop));
@@ -2774,6 +2784,16 @@ public Response<Long> zrevrank(byte[] key, byte[] member) {
     return appendCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public Response<KeyValue<Long, Double>> zrankWithScore(byte[] key, byte[] member) {
+    return appendCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public Response<KeyValue<Long, Double>> zrevrankWithScore(byte[] key, byte[] member) {
+    return appendCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public Response<List<byte[]>> zrange(byte[] key, long start, long stop) {
     return appendCommand(commandObjects.zrange(key, start, stop));

diff --git a/src/main/java/redis/clients/jedis/UnifiedJedis.java b/src/main/java/redis/clients/jedis/UnifiedJedis.java
@@ -1789,6 +1789,16 @@ public Long zrevrank(String key, String member) {
     return executeCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public KeyValue<Long, Double> zrankWithScore(String key, String member) {
+    return executeCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public KeyValue<Long, Double> zrevrankWithScore(String key, String member) {
+    return executeCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public long zrem(byte[] key, byte[]... members) {
     return executeCommand(commandObjects.zrem(key, members));
@@ -1814,6 +1824,16 @@ public Long zrevrank(byte[] key, byte[] member) {
     return executeCommand(commandObjects.zrevrank(key, member));
   }
 
+  @Override
+  public KeyValue<Long, Double> zrankWithScore(byte[] key, byte[] member) {
+    return executeCommand(commandObjects.zrankWithScore(key, member));
+  }
+
+  @Override
+  public KeyValue<Long, Double> zrevrankWithScore(byte[] key, byte[] member) {
+    return executeCommand(commandObjects.zrevrankWithScore(key, member));
+  }
+
   @Override
   public String zrandmember(String key) {
     return executeCommand(commandObjects.zrandmember(key));

diff --git a/src/main/java/redis/clients/jedis/commands/SortedSetBinaryCommands.java b/src/main/java/redis/clients/jedis/commands/SortedSetBinaryCommands.java
@@ -32,6 +32,10 @@ public interface SortedSetBinaryCommands {
 
   Long zrevrank(byte[] key, byte[] member);
 
+  KeyValue<Long, Double> zrankWithScore(byte[] key, byte[] member);
+
+  KeyValue<Long, Double> zrevrankWithScore(byte[] key, byte[] member);
+
   List<byte[]> zrange(byte[] key, long start, long stop);
 
   List<byte[]> zrevrange(byte[] key, long start, long stop);

diff --git a/src/main/java/redis/clients/jedis/commands/SortedSetCommands.java b/src/main/java/redis/clients/jedis/commands/SortedSetCommands.java
@@ -153,6 +153,24 @@ public interface SortedSetCommands {
    */
   Long zrevrank(String key, String member);
 
+  /**
+   * Returns the rank and the score of member in the sorted set stored at key, with the scores
+   * ordered from low to high.
+   * @param key the key
+   * @param member the member
+   * @return the KeyValue contains rank and score.
+   */
+  KeyValue<Long, Double> zrankWithScore(String key, String member);
+
+  /**
+   * Returns the rank and the score of member in the sorted set stored at key, with the scores
+   * ordered from high to low.
+   * @param key the key
+   * @param member the member
+   * @return the KeyValue contains rank and score.
+   */
+  KeyValue<Long, Double> zrevrankWithScore(String key, String member);
+
   /**
    * Returns the specified range of elements in the sorted set stored at key.
    * <p>

diff --git a/src/main/java/redis/clients/jedis/commands/SortedSetPipelineBinaryCommands.java b/src/main/java/redis/clients/jedis/commands/SortedSetPipelineBinaryCommands.java
@@ -33,6 +33,10 @@ public interface SortedSetPipelineBinaryCommands {
 
   Response<Long> zrevrank(byte[] key, byte[] member);
 
+  Response<KeyValue<Long, Double>> zrankWithScore(byte[] key, byte[] member);
+
+  Response<KeyValue<Long, Double>> zrevrankWithScore(byte[] key, byte[] member);
+
   Response<List<byte[]>> zrange(byte[] key, long start, long stop);
 
   Response<List<byte[]>> zrevrange(byte[] key, long start, long stop);

diff --git a/src/main/java/redis/clients/jedis/commands/SortedSetPipelineCommands.java b/src/main/java/redis/clients/jedis/commands/SortedSetPipelineCommands.java
@@ -34,6 +34,10 @@ public interface SortedSetPipelineCommands {
 
   Response<Long> zrevrank(String key, String member);
 
+  Response<KeyValue<Long, Double>> zrankWithScore(String key, String member);
+
+  Response<KeyValue<Long, Double>> zrevrankWithScore(String key, String member);
+
   Response<List<String>> zrange(String key, long start, long stop);
 
   Response<List<String>> zrevrange(String key, long start, long stop);

diff --git a/src/test/java/redis/clients/jedis/commands/jedis/SortedSetCommandsTest.java b/src/test/java/redis/clients/jedis/commands/jedis/SortedSetCommandsTest.java
@@ -402,6 +402,36 @@ public void zrank() {
     assertNull(jedis.zrank(bcar, bb));
   }
 
+  @Test
+  public void zrankWithScore() {
+    jedis.zadd("foo", 1d, "a");
+    jedis.zadd("foo", 2d, "b");
+
+    KeyValue<Long, Double> keyValue = jedis.zrankWithScore("foo", "a");
+    assertEquals(Long.valueOf(0), keyValue.getKey());
+    assertEquals(Double.valueOf(1d), keyValue.getValue());
+
+    keyValue = jedis.zrankWithScore("foo", "b");
+    assertEquals(Long.valueOf(1), keyValue.getKey());
+    assertEquals(Double.valueOf(2d), keyValue.getValue());
+
+    assertNull(jedis.zrankWithScore("car", "b"));
+
+    // Binary
+    jedis.zadd(bfoo, 1d, ba);
+    jedis.zadd(bfoo, 2d, bb);
+
+    keyValue = jedis.zrankWithScore(bfoo, ba);
+    assertEquals(Long.valueOf(0), keyValue.getKey());
+    assertEquals(Double.valueOf(1d), keyValue.getValue());
+
+    keyValue = jedis.zrankWithScore(bfoo, bb);
+    assertEquals(Long.valueOf(1), keyValue.getKey());
+    assertEquals(Double.valueOf(2d), keyValue.getValue());
+
+    assertNull(jedis.zrankWithScore(bcar, bb));
+  }
+
   @Test
   public void zrevrank() {
     jedis.zadd("foo", 1d, "a");

diff --git a/src/test/java/redis/clients/jedis/commands/unified/SortedSetCommandsTestBase.java b/src/test/java/redis/clients/jedis/commands/unified/SortedSetCommandsTestBase.java
@@ -459,6 +459,36 @@ public void zrank() {
     assertNull(jedis.zrank(bcar, bb));
   }
 
+  @Test
+  public void zrankWithScore() {
+    jedis.zadd("foo", 1d, "a");
+    jedis.zadd("foo", 2d, "b");
+
+    KeyValue<Long, Double> keyValue = jedis.zrankWithScore("foo", "a");
+    assertEquals(Long.valueOf(0), keyValue.getKey());
+    assertEquals(Double.valueOf(1d), keyValue.getValue());
+
+    keyValue = jedis.zrankWithScore("foo", "b");
+    assertEquals(Long.valueOf(1), keyValue.getKey());
+    assertEquals(Double.valueOf(2d), keyValue.getValue());
+
+    assertNull(jedis.zrankWithScore("car", "b"));
+
+    // Binary
+    jedis.zadd(bfoo, 1d, ba);
+    jedis.zadd(bfoo, 2d, bb);
+
+    keyValue = jedis.zrankWithScore(bfoo, ba);
+    assertEquals(Long.valueOf(0), keyValue.getKey());
+    assertEquals(Double.valueOf(1d), keyValue.getValue());
+
+    keyValue = jedis.zrankWithScore(bfoo, bb);
+    assertEquals(Long.valueOf(1), keyValue.getKey());
+    assertEquals(Double.valueOf(2d), keyValue.getValue());
+
+    assertNull(jedis.zrankWithScore(bcar, bb));
+  }
+
   @Test
   public void zrevrank() {
     jedis.zadd("foo", 1d, "a");