Support TDIGEST.[BY][REV]RANK commands (#3158)

redis · Oct 6, 2022 · ffff3ec · ffff3ec
1 parent 4a06e11
commit ffff3ec
Show file tree

Hide file tree

Showing 9 changed files with 149 additions and 16 deletions.
diff --git a/src/main/java/redis/clients/jedis/CommandObjects.java b/src/main/java/redis/clients/jedis/CommandObjects.java
@@ -3907,27 +3907,18 @@ public final CommandObject<Map<String, Object>> tdigestInfo(String key) {
   }
 
   public final CommandObject<String> tdigestAdd(String key, double... values) {
-    CommandArguments args = commandArguments(TDigestCommand.ADD).key(key);
-    for (double value : values) {
-      args.add(value);
-    }
-    return new CommandObject<>(args, BuilderFactory.STRING);
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.ADD).key(key), values),
+        BuilderFactory.STRING);
   }
 
   public final CommandObject<List<Double>> tdigestCDF(String key, double... values) {
-    CommandArguments args = commandArguments(TDigestCommand.CDF).key(key);
-    for (double value : values) {
-      args.add(value);
-    }
-    return new CommandObject<>(args, BuilderFactory.DOUBLE_LIST);
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.CDF).key(key), values),
+        BuilderFactory.DOUBLE_LIST);
   }
 
   public final CommandObject<List<Double>> tdigestQuantile(String key, double... quantiles) {
-    CommandArguments args = commandArguments(TDigestCommand.QUANTILE).key(key);
-    for (double quantile : quantiles) {
-      args.add(quantile);
-    }
-    return new CommandObject<>(args, BuilderFactory.DOUBLE_LIST);
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.QUANTILE).key(key),
+        quantiles), BuilderFactory.DOUBLE_LIST);
   }
 
   public final CommandObject<Double> tdigestMin(String key) {
@@ -3942,6 +3933,26 @@ public final CommandObject<Double> tdigestTrimmedMean(String key, double lowCutQ
     return new CommandObject<>(commandArguments(TDigestCommand.TRIMMED_MEAN).key(key).add(lowCutQuantile)
         .add(highCutQuantile), BuilderFactory.DOUBLE);
   }
+
+  public final CommandObject<List<Long>> tdigestRank(String key, double... values) {
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.RANK).key(key),
+        values), BuilderFactory.LONG_LIST);
+  }
+
+  public final CommandObject<List<Long>> tdigestRevRank(String key, double... values) {
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.REVRANK).key(key),
+        values), BuilderFactory.LONG_LIST);
+  }
+
+  public final CommandObject<List<Double>> tdigestByRank(String key, long... ranks) {
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.BYRANK).key(key),
+        ranks), BuilderFactory.DOUBLE_LIST);
+  }
+
+  public final CommandObject<List<Double>> tdigestByRevRank(String key, long... ranks) {
+    return new CommandObject<>(addFlatArgs(commandArguments(TDigestCommand.BYREVRANK).key(key),
+        ranks), BuilderFactory.DOUBLE_LIST);
+  }
   // RedisBloom commands
 
   // RedisGraph commands
@@ -4004,6 +4015,20 @@ public List<T> build(Object data) {
     }
   }
 
+  private CommandArguments addFlatArgs(CommandArguments args, long... values) {
+    for (long value : values) {
+      args.add(value);
+    }
+    return args;
+  }
+
+  private CommandArguments addFlatArgs(CommandArguments args, double... values) {
+    for (double value : values) {
+      args.add(value);
+    }
+    return args;
+  }
+
   private CommandArguments addFlatKeyValueArgs(CommandArguments args, String... keyvalues) {
     for (int i = 0; i < keyvalues.length; i += 2) {
       args.key(keyvalues[i]).add(keyvalues[i + 1]);

diff --git a/src/main/java/redis/clients/jedis/MultiNodePipelineBase.java b/src/main/java/redis/clients/jedis/MultiNodePipelineBase.java
@@ -4209,6 +4209,26 @@ public Response<Double> tdigestMax(String key) {
   public Response<Double> tdigestTrimmedMean(String key, double lowCutQuantile, double highCutQuantile) {
     return appendCommand(commandObjects.tdigestTrimmedMean(key, lowCutQuantile, highCutQuantile));
   }
+
+  @Override
+  public Response<List<Long>> tdigestRank(String key, double... values) {
+    return appendCommand(commandObjects.tdigestRank(key, values));
+  }
+
+  @Override
+  public Response<List<Long>> tdigestRevRank(String key, double... values) {
+    return appendCommand(commandObjects.tdigestRevRank(key, values));
+  }
+
+  @Override
+  public Response<List<Double>> tdigestByRank(String key, long... ranks) {
+    return appendCommand(commandObjects.tdigestByRank(key, ranks));
+  }
+
+  @Override
+  public Response<List<Double>> tdigestByRevRank(String key, long... ranks) {
+    return appendCommand(commandObjects.tdigestByRevRank(key, ranks));
+  }
   // RedisBloom commands
 
   // RedisGraph commands

diff --git a/src/main/java/redis/clients/jedis/Pipeline.java b/src/main/java/redis/clients/jedis/Pipeline.java
@@ -4188,6 +4188,26 @@ public Response<Double> tdigestMax(String key) {
   public Response<Double> tdigestTrimmedMean(String key, double lowCutQuantile, double highCutQuantile) {
     return appendCommand(commandObjects.tdigestTrimmedMean(key, lowCutQuantile, highCutQuantile));
   }
+
+  @Override
+  public Response<List<Long>> tdigestRank(String key, double... values) {
+    return appendCommand(commandObjects.tdigestRank(key, values));
+  }
+
+  @Override
+  public Response<List<Long>> tdigestRevRank(String key, double... values) {
+    return appendCommand(commandObjects.tdigestRevRank(key, values));
+  }
+
+  @Override
+  public Response<List<Double>> tdigestByRank(String key, long... ranks) {
+    return appendCommand(commandObjects.tdigestByRank(key, ranks));
+  }
+
+  @Override
+  public Response<List<Double>> tdigestByRevRank(String key, long... ranks) {
+    return appendCommand(commandObjects.tdigestByRevRank(key, ranks));
+  }
   // RedisBloom commands
 
   // RedisGraph commands

diff --git a/src/main/java/redis/clients/jedis/TransactionBase.java b/src/main/java/redis/clients/jedis/TransactionBase.java
@@ -4256,6 +4256,26 @@ public Response<Double> tdigestMax(String key) {
   public Response<Double> tdigestTrimmedMean(String key, double lowCutQuantile, double highCutQuantile) {
     return appendCommand(commandObjects.tdigestTrimmedMean(key, lowCutQuantile, highCutQuantile));
   }
+
+  @Override
+  public Response<List<Long>> tdigestRank(String key, double... values) {
+    return appendCommand(commandObjects.tdigestRank(key, values));
+  }
+
+  @Override
+  public Response<List<Long>> tdigestRevRank(String key, double... values) {
+    return appendCommand(commandObjects.tdigestRevRank(key, values));
+  }
+
+  @Override
+  public Response<List<Double>> tdigestByRank(String key, long... ranks) {
+    return appendCommand(commandObjects.tdigestByRank(key, ranks));
+  }
+
+  @Override
+  public Response<List<Double>> tdigestByRevRank(String key, long... ranks) {
+    return appendCommand(commandObjects.tdigestByRevRank(key, ranks));
+  }
   // RedisBloom commands
 
   // RedisGraph commands

diff --git a/src/main/java/redis/clients/jedis/UnifiedJedis.java b/src/main/java/redis/clients/jedis/UnifiedJedis.java
@@ -4401,6 +4401,26 @@ public double tdigestMax(String key) {
   public double tdigestTrimmedMean(String key, double lowCutQuantile, double highCutQuantile) {
     return executeCommand(commandObjects.tdigestTrimmedMean(key, lowCutQuantile, highCutQuantile));
   }
+
+  @Override
+  public List<Long> tdigestRank(String key, double... values) {
+    return executeCommand(commandObjects.tdigestRank(key, values));
+  }
+
+  @Override
+  public List<Long> tdigestRevRank(String key, double... values) {
+    return executeCommand(commandObjects.tdigestRevRank(key, values));
+  }
+
+  @Override
+  public List<Double> tdigestByRank(String key, long... ranks) {
+    return executeCommand(commandObjects.tdigestByRank(key, ranks));
+  }
+
+  @Override
+  public List<Double> tdigestByRevRank(String key, long... ranks) {
+    return executeCommand(commandObjects.tdigestByRevRank(key, ranks));
+  }
   // RedisBloom commands
 
   // RedisGraph commands

diff --git a/src/main/java/redis/clients/jedis/bloom/RedisBloomProtocol.java b/src/main/java/redis/clients/jedis/bloom/RedisBloomProtocol.java
@@ -102,7 +102,8 @@ public byte[] getRaw() {
 
   public enum TDigestCommand implements ProtocolCommand {
 
-    CREATE, INFO, ADD, RESET, MERGE, CDF, QUANTILE, MIN, MAX, TRIMMED_MEAN;
+    CREATE, INFO, ADD, RESET, MERGE, CDF, QUANTILE, MIN, MAX, TRIMMED_MEAN,
+    RANK, REVRANK, BYRANK, BYREVRANK;
 
     private final byte[] raw;
 

diff --git a/src/main/java/redis/clients/jedis/bloom/commands/TDigestSketchCommands.java b/src/main/java/redis/clients/jedis/bloom/commands/TDigestSketchCommands.java
@@ -111,4 +111,12 @@ public interface TDigestSketchCommands {
    * @return estimation of the mean value
    */
   double tdigestTrimmedMean(String key, double lowCutQuantile, double highCutQuantile);
+
+  List<Long> tdigestRank(String key, double... values);
+
+  List<Long> tdigestRevRank(String key, double... values);
+
+  List<Double> tdigestByRank(String key, long... ranks);
+
+  List<Double> tdigestByRevRank(String key, long... ranks);
 }
diff --git a/src/main/java/redis/clients/jedis/bloom/commands/TDigestSketchPipelineCommands.java b/src/main/java/redis/clients/jedis/bloom/commands/TDigestSketchPipelineCommands.java
@@ -30,4 +30,12 @@ public interface TDigestSketchPipelineCommands {
   Response<Double> tdigestMax(String key);
 
   Response<Double> tdigestTrimmedMean(String key, double lowCutQuantile, double highCutQuantile);
+
+  Response<List<Long>> tdigestRank(String key, double... values);
+
+  Response<List<Long>> tdigestRevRank(String key, double... values);
+
+  Response<List<Double>> tdigestByRank(String key, long... ranks);
+
+  Response<List<Double>> tdigestByRevRank(String key, long... ranks);
 }
diff --git a/src/test/java/redis/clients/jedis/modules/bloom/TDigestTest.java b/src/test/java/redis/clients/jedis/modules/bloom/TDigestTest.java
@@ -177,6 +177,17 @@ public void trimmedMean() {
     assertEquals(4.5, client.tdigestTrimmedMean(key, 0.0, 0.5), 0.01);
     assertEquals(14.5, client.tdigestTrimmedMean(key, 0.5, 1.0), 0.01);
   }
+
+  @Test
+  public void rankCommands() {
+    final String key = "ranks";
+    client.tdigestCreate(key);
+    client.tdigestAdd(key, 2d, 3d, 5d);
+    assertEquals(Arrays.asList(1l, 2l), client.tdigestRank(key, 2d, 4d));
+    assertEquals(Arrays.asList(0l, 1l), client.tdigestRevRank(key, 5d, 4d));
+    assertEquals(Arrays.asList(2d, 3d), client.tdigestByRank(key, 0l, 1l));
+    assertEquals(Arrays.asList(5d, 3d), client.tdigestByRevRank(key, 1l, 2l));
+  }
 //
 //  private static KeyValue<Double, Long> randomValueWeight() {
 //    return new KeyValue<>(random.nextDouble() * 10000, Math.abs(random.nextInt()) + 1l);