Java: Add PFCOUNT command. (#1222)

* Add `PFCOUNT` command. (#167) Signed-off-by: Yury-Fridlyand <[email protected]>
valkey-io · Apr 5, 2024 · 63ab900 · 63ab900
1 parent 334c864
commit 63ab900
Show file tree

Hide file tree

Showing 9 changed files with 97 additions and 0 deletions.
diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto
@@ -137,6 +137,7 @@ enum RequestType {
     Brpop = 93;
     Hkeys = 94;
     PfAdd = 96;
+    PfCount = 97;
     Blpop = 100;
     RPushX = 102;
     LPushX = 103;

diff --git a/glide-core/src/socket_listener.rs b/glide-core/src/socket_listener.rs
@@ -364,6 +364,7 @@ fn get_command(request: &Command) -> Option<Cmd> {
         RequestType::Brpop => Some(cmd("BRPOP")),
         RequestType::Hkeys => Some(cmd("HKEYS")),
         RequestType::PfAdd => Some(cmd("PFADD")),
+        RequestType::PfCount => Some(cmd("PFCOUNT")),
         RequestType::RPushX => Some(cmd("RPUSHX")),
         RequestType::LPushX => Some(cmd("LPUSHX")),
         RequestType::Blpop => Some(cmd("BLPOP")),

diff --git a/java/client/src/main/java/glide/api/BaseClient.java b/java/client/src/main/java/glide/api/BaseClient.java
@@ -43,6 +43,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
 import static redis_request.RedisRequestOuterClass.RequestType.Persist;
 import static redis_request.RedisRequestOuterClass.RequestType.PfAdd;
+import static redis_request.RedisRequestOuterClass.RequestType.PfCount;
 import static redis_request.RedisRequestOuterClass.RequestType.RPop;
 import static redis_request.RedisRequestOuterClass.RequestType.RPush;
 import static redis_request.RedisRequestOuterClass.RequestType.RPushX;
@@ -789,4 +790,9 @@ public CompletableFuture<Long> pfadd(@NonNull String key, @NonNull String[] elem
         String[] arguments = ArrayUtils.addFirst(elements, key);
         return commandManager.submitNewCommand(PfAdd, arguments, this::handleLongResponse);
     }
+
+    @Override
+    public CompletableFuture<Long> pfcount(@NonNull String[] keys) {
+        return commandManager.submitNewCommand(PfCount, keys, this::handleLongResponse);
+    }
 }
diff --git a/java/client/src/main/java/glide/api/commands/HyperLogLogBaseCommands.java b/java/client/src/main/java/glide/api/commands/HyperLogLogBaseCommands.java
@@ -35,4 +35,20 @@ public interface HyperLogLogBaseCommands {
      * }</pre>
      */
     CompletableFuture<Long> pfadd(String key, String[] elements);
+
+    /**
+     * Estimates the cardinality of the data stored in a HyperLogLog structure for a single key or
+     * calculates the combined cardinality of multiple keys by merging their HyperLogLogs temporarily.
+     *
+     * @see <a href="https://redis.io/commands/pfcount/">redis.io</a> for details.
+     * @param keys The keys of the HyperLogLog data structures to be analyzed.
+     * @return The approximated cardinality of given HyperLogLog data structures.<br>
+     *     The cardinality of a key that does not exist is <code>0</code>.
+     * @example
+     *     <pre>{@code
+     * Long result = client.pfcount("hll_1", "hll_2").get();
+     * assert result == 42L; // Count of unique elements in multiple data structures
+     * }</pre>
+     */
+    CompletableFuture<Long> pfcount(String[] keys);
 }
diff --git a/java/client/src/main/java/glide/api/models/BaseTransaction.java b/java/client/src/main/java/glide/api/models/BaseTransaction.java
@@ -52,6 +52,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
 import static redis_request.RedisRequestOuterClass.RequestType.Persist;
 import static redis_request.RedisRequestOuterClass.RequestType.PfAdd;
+import static redis_request.RedisRequestOuterClass.RequestType.PfCount;
 import static redis_request.RedisRequestOuterClass.RequestType.Ping;
 import static redis_request.RedisRequestOuterClass.RequestType.RPop;
 import static redis_request.RedisRequestOuterClass.RequestType.RPush;
@@ -1757,6 +1758,22 @@ public T pfadd(@NonNull String key, @NonNull String[] elements) {
         return getThis();
     }
 
+    /**
+     * Estimates the cardinality of the data stored in a HyperLogLog structure for a single key or
+     * calculates the combined cardinality of multiple keys by merging their HyperLogLogs temporarily.
+     *
+     * @see <a href="https://redis.io/commands/pfcount/">redis.io</a> for details.
+     * @param keys The keys of the HyperLogLog data structures to be analyzed.
+     * @return Command Response - The approximated cardinality of given HyperLogLog data structures.
+     *     <br>
+     *     The cardinality of a key that does not exist is <code>0</code>.
+     */
+    public T pfcount(@NonNull String[] keys) {
+        ArgsArray commandArgs = buildArgs(keys);
+        protobufTransaction.addCommands(buildCommand(PfCount, commandArgs));
+        return getThis();
+    }
+
     /** Build protobuf {@link Command} object for given command and arguments. */
     protected Command buildCommand(RequestType requestType) {
         return buildCommand(requestType, buildArgs());

diff --git a/java/client/src/test/java/glide/api/RedisClientTest.java b/java/client/src/test/java/glide/api/RedisClientTest.java
@@ -71,6 +71,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
 import static redis_request.RedisRequestOuterClass.RequestType.Persist;
 import static redis_request.RedisRequestOuterClass.RequestType.PfAdd;
+import static redis_request.RedisRequestOuterClass.RequestType.PfCount;
 import static redis_request.RedisRequestOuterClass.RequestType.Ping;
 import static redis_request.RedisRequestOuterClass.RequestType.RPop;
 import static redis_request.RedisRequestOuterClass.RequestType.RPush;
@@ -2555,4 +2556,28 @@ public void pfadd_returns_success() {
         assertEquals(testResponse, response);
         assertEquals(value, payload);
     }
+
+    @SneakyThrows
+    @Test
+    public void pfcount_returns_success() {
+        // setup
+        String[] keys = new String[] {"a", "b", "c"};
+        Long value = 1L;
+
+        CompletableFuture<Long> testResponse = new CompletableFuture<>();
+        testResponse.complete(value);
+
+        // match on protobuf request
+        when(commandManager.<Long>submitNewCommand(eq(PfCount), eq(keys), any()))
+                .thenReturn(testResponse);
+
+        // exercise
+        CompletableFuture<Long> response = service.pfcount(keys);
+        Long payload = response.get();
+
+        // verify
+        assertEquals(testResponse, response);
+        assertEquals(value, payload);
+        assertEquals(payload, response.get());
+    }
 }
diff --git a/java/client/src/test/java/glide/api/models/TransactionTests.java b/java/client/src/test/java/glide/api/models/TransactionTests.java
@@ -50,6 +50,7 @@
 import static redis_request.RedisRequestOuterClass.RequestType.PTTL;
 import static redis_request.RedisRequestOuterClass.RequestType.Persist;
 import static redis_request.RedisRequestOuterClass.RequestType.PfAdd;
+import static redis_request.RedisRequestOuterClass.RequestType.PfCount;
 import static redis_request.RedisRequestOuterClass.RequestType.Ping;
 import static redis_request.RedisRequestOuterClass.RequestType.RPop;
 import static redis_request.RedisRequestOuterClass.RequestType.RPush;
@@ -544,6 +545,9 @@ InfScoreBound.NEGATIVE_INFINITY, new ScoreBoundary(3, false), new Limit(1, 2)),
                         PfAdd,
                         ArgsArray.newBuilder().addArgs("hll").addArgs("a").addArgs("b").addArgs("c").build()));
 
+        transaction.pfcount(new String[] {"hll1", "hll2"});
+        results.add(Pair.of(PfCount, ArgsArray.newBuilder().addArgs("hll1").addArgs("hll2").build()));
+
         var protobufTransaction = transaction.getProtobufTransaction().build();
 
         for (int idx = 0; idx < protobufTransaction.getCommandsCount(); idx++) {

diff --git a/java/integTest/src/test/java/glide/SharedCommandTests.java b/java/integTest/src/test/java/glide/SharedCommandTests.java
@@ -1641,4 +1641,28 @@ public void pfadd(BaseClient client) {
                 assertThrows(ExecutionException.class, () -> client.pfadd("foo", new String[0]).get());
         assertTrue(executionException.getCause() instanceof RequestException);
     }
+
+    @SneakyThrows
+    @ParameterizedTest
+    @MethodSource("getClients")
+    public void pfcount(BaseClient client) {
+        String key1 = "{test}-hll1-" + UUID.randomUUID();
+        String key2 = "{test}-hll2-" + UUID.randomUUID();
+        String key3 = "{test}-hll3-" + UUID.randomUUID();
+        assertEquals(1, client.pfadd(key1, new String[] {"a", "b", "c"}).get());
+        assertEquals(1, client.pfadd(key2, new String[] {"b", "c", "d"}).get());
+        assertEquals(3, client.pfcount(new String[] {key1}).get());
+        assertEquals(3, client.pfcount(new String[] {key2}).get());
+        assertEquals(4, client.pfcount(new String[] {key1, key2}).get());
+        assertEquals(4, client.pfcount(new String[] {key1, key2, key3}).get());
+        // empty HyperLogLog data set
+        assertEquals(1, client.pfadd(key3, new String[0]).get());
+        assertEquals(0, client.pfcount(new String[] {key3}).get());
+
+        // Key exists, but it is not a HyperLogLog
+        assertEquals(OK, client.set("foo", "bar").get());
+        ExecutionException executionException =
+                assertThrows(ExecutionException.class, () -> client.pfcount(new String[] {"foo"}).get());
+        assertTrue(executionException.getCause() instanceof RequestException);
+    }
 }
diff --git a/java/integTest/src/test/java/glide/TransactionTestUtilities.java b/java/integTest/src/test/java/glide/TransactionTestUtilities.java
@@ -23,6 +23,7 @@ public class TransactionTestUtilities {
     private static final String key8 = "{key}" + UUID.randomUUID();
     private static final String key9 = "{key}" + UUID.randomUUID();
     private static final String hllKey1 = "{key}:hllKey1-" + UUID.randomUUID();
+    private static final String hllKey2 = "{key}:hllKey2-" + UUID.randomUUID();
     private static final String value1 = UUID.randomUUID().toString();
     private static final String value2 = UUID.randomUUID().toString();
     private static final String value3 = UUID.randomUUID().toString();
@@ -126,6 +127,7 @@ public static BaseTransaction<?> transactionTest(BaseTransaction<?> baseTransact
                 .brpop(new String[] {listKey3}, 0.01);
 
         baseTransaction.pfadd(hllKey1, new String[] {"a", "b", "c"});
+        baseTransaction.pfcount(new String[] {hllKey1, hllKey2});
 
         return baseTransaction;
     }
@@ -203,6 +205,7 @@ public static Object[] transactionTestResult() {
             new String[] {listKey3, value3}, // blpop(new String[] { listKey3 }, 0.01)
             new String[] {listKey3, value1}, // brpop(new String[] { listKey3 }, 0.01);
             1L, // pfadd(hllKey1, new String[] {"a", "b", "c"})
+            3L, // pfcount(new String[] { hllKey1, hllKey2 });
         };
     }
 }