Support TOPK.LIST with WITHCOUNT option

sazzad16 · Aug 8, 2023 · 69dd3dc · 69dd3dc
1 parent 2cd2601
commit 69dd3dc
Show file tree

Hide file tree

Showing 9 changed files with 90 additions and 7 deletions.
diff --git a/src/main/java/redis/clients/jedis/BuilderFactory.java b/src/main/java/redis/clients/jedis/BuilderFactory.java
@@ -415,6 +415,37 @@ public String toString() {
     }
   };
 
+  public static final Builder<Map<String, Long>> STRING_LONG_MAP = new Builder<Map<String, Long>>() {
+    @Override
+    @SuppressWarnings("unchecked")
+    public Map<String, Long> build(Object data) {
+      final List<Object> list = (List<Object>) data;
+      if (list.isEmpty()) return Collections.emptyMap();
+
+      if (list.get(0) instanceof KeyValue) {
+        final Map<String, Long> map = new LinkedHashMap<>(list.size(), 1f);
+        final Iterator iterator = list.iterator();
+        while (iterator.hasNext()) {
+          KeyValue kv = (KeyValue) iterator.next();
+          map.put(STRING.build(kv.getKey()), LONG.build(kv.getValue()));
+        }
+        return map;
+      } else {
+        final Map<String, Long> map = new LinkedHashMap<>(list.size() / 2, 1f);
+        final Iterator iterator = list.iterator();
+        while (iterator.hasNext()) {
+          map.put(STRING.build(iterator.next()), LONG.build(iterator.next()));
+        }
+        return map;
+      }
+    }
+
+    @Override
+    public String toString() {
+      return "Map<String, String>";
+    }
+  };
+
   public static final Builder<KeyValue<String, String>> KEYED_ELEMENT = new Builder<KeyValue<String, String>>() {
     @Override
     @SuppressWarnings("unchecked")

diff --git a/src/main/java/redis/clients/jedis/CommandObjects.java b/src/main/java/redis/clients/jedis/CommandObjects.java
@@ -4056,6 +4056,11 @@ public final CommandObject<List<String>> topkList(String key) {
     return new CommandObject<>(commandArguments(TopKCommand.LIST).key(key), BuilderFactory.STRING_LIST);
   }
 
+  public final CommandObject<Map<String, Long>> topkListWithCount(String key) {
+    return new CommandObject<>(commandArguments(TopKCommand.LIST).key(key)
+        .add(RedisBloomKeyword.WITHCOUNT), BuilderFactory.STRING_LONG_MAP);
+  }
+
   public final CommandObject<Map<String, Object>> topkInfo(String key) {
     return new CommandObject<>(commandArguments(TopKCommand.INFO).key(key), BuilderFactory.ENCODED_OBJECT_MAP);
   }

diff --git a/src/main/java/redis/clients/jedis/PipelineBase.java b/src/main/java/redis/clients/jedis/PipelineBase.java
@@ -4059,6 +4059,11 @@ public Response<List<String>> topkList(String key) {
     return appendCommand(commandObjects.topkList(key));
   }
 
+  @Override
+  public Response<Map<String, Long>> topkListWithCount(String key) {
+    return appendCommand(commandObjects.topkListWithCount(key));
+  }
+
   @Override
   public Response<Map<String, Object>> topkInfo(String key) {
     return appendCommand(commandObjects.topkInfo(key));

diff --git a/src/main/java/redis/clients/jedis/TransactionBase.java b/src/main/java/redis/clients/jedis/TransactionBase.java
@@ -4227,6 +4227,11 @@ public Response<List<String>> topkList(String key) {
     return appendCommand(commandObjects.topkList(key));
   }
 
+  @Override
+  public Response<Map<String, Long>> topkListWithCount(String key) {
+    return appendCommand(commandObjects.topkListWithCount(key));
+  }
+
   @Override
   public Response<Map<String, Object>> topkInfo(String key) {
     return appendCommand(commandObjects.topkInfo(key));

diff --git a/src/main/java/redis/clients/jedis/UnifiedJedis.java b/src/main/java/redis/clients/jedis/UnifiedJedis.java
@@ -4583,6 +4583,11 @@ public List<String> topkList(String key) {
     return executeCommand(commandObjects.topkList(key));
   }
 
+  @Override
+  public Map<String, Long> topkListWithCount(String key) {
+    return executeCommand(commandObjects.topkListWithCount(key));
+  }
+
   @Override
   public Map<String, Object> topkInfo(String key) {
     return executeCommand(commandObjects.topkInfo(key));

diff --git a/src/main/java/redis/clients/jedis/bloom/RedisBloomProtocol.java b/src/main/java/redis/clients/jedis/bloom/RedisBloomProtocol.java
@@ -120,7 +120,7 @@ public byte[] getRaw() {
   public enum RedisBloomKeyword implements Rawable {
 
     CAPACITY, ERROR, NOCREATE, EXPANSION, NONSCALING, BUCKETSIZE, MAXITERATIONS, ITEMS, WEIGHTS,
-    COMPRESSION, OVERRIDE;
+    COMPRESSION, OVERRIDE, WITHCOUNT;
 
     private final byte[] raw;
 

diff --git a/src/main/java/redis/clients/jedis/bloom/commands/TopKFilterCommands.java b/src/main/java/redis/clients/jedis/bloom/commands/TopKFilterCommands.java
@@ -1,5 +1,6 @@
 package redis.clients.jedis.bloom.commands;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -35,6 +36,18 @@ public interface TopKFilterCommands {
    */
   List<String> topkAdd(String key, String... items);
 
+  /**
+   * {@code TOPK.INCRBY {key} {item} {increment}}
+   *
+   * @param key
+   * @param item
+   * @param increment
+   * @return item dropped from list
+   */
+  default String topkIncrBy(String key, String item, long increment) {
+    return topkIncrBy(key, Collections.singletonMap(item, increment)).get(0);
+  }
+
   /**
    * {@code TOPK.INCRBY {key} {item} {increment} [{item} {increment} ...]}
    *
@@ -61,6 +74,14 @@ public interface TopKFilterCommands {
    */
   List<String> topkList(String key);
 
+  /**
+   * {@code TOPK.LIST {key} WITHCOUNT}
+   *
+   * @param key
+   * @return k (or less) items in Top K list
+   */
+  Map<String, Long> topkListWithCount(String key);
+
   /**
    * {@code TOPK.INFO {key}}
    *

diff --git a/src/main/java/redis/clients/jedis/bloom/commands/TopKFilterPipelineCommands.java b/src/main/java/redis/clients/jedis/bloom/commands/TopKFilterPipelineCommands.java
@@ -18,5 +18,7 @@ public interface TopKFilterPipelineCommands {
 
   Response<List<String>> topkList(String key);
 
+  Response<Map<String, Long>> topkListWithCount(String key);
+
   Response<Map<String, Object>> topkInfo(String key);
 }
diff --git a/src/test/java/redis/clients/jedis/modules/bloom/TopKTest.java b/src/test/java/redis/clients/jedis/modules/bloom/TopKTest.java
@@ -1,10 +1,11 @@
 package redis.clients.jedis.modules.bloom;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.TreeSet;
+import java.util.Map;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -30,12 +31,20 @@ public void createTopKFilter() {
 
     assertEquals(Arrays.asList(true, false, true), client.topkQuery("aaa", "bb", "gg", "cc"));
 
-    assertEquals(new TreeSet<>(Arrays.asList("bb", "cc")), new TreeSet<>(client.topkList("aaa")));
+    assertEquals(Arrays.asList("bb", "cc"), client.topkList("aaa"));
 
-//    assertEquals(null, client.topkIncrBy("aaa", "ff", 10));
-    assertEquals(Collections.<String>singletonList(null),
-        client.topkIncrBy("aaa", Collections.singletonMap("ff", 10L)));
+    Map<String, Long> listWithCount = client.topkListWithCount("aaa");
+    assertEquals(2, listWithCount.size());
+    listWithCount.forEach((item, count) -> {
+      assertTrue(Arrays.asList("bb", "cc").contains(item));
+      assertEquals(Long.valueOf(1), count);
+    });
+
+    assertEquals(null, client.topkIncrBy("aaa", "ff", 5));
+    assertEquals(Arrays.asList("ff", "bb", "cc"), client.topkList("aaa"));
 
-    assertEquals(new TreeSet<>(Arrays.asList("bb", "cc", "ff")), new TreeSet<>(client.topkList("aaa")));
+    assertEquals(Collections.<String>singletonList(null),
+        client.topkIncrBy("aaa", Collections.singletonMap("ff", 8L)));
+    assertEquals(Long.valueOf(13), client.topkListWithCount("aaa").get("ff"));
   }
 }