Python: add PFMERGE command (valkey-io#1497)

Python: add PFMERGE command (#322)

Author: aaron-congo

CHANGELOG.md

@@ -6,6 +6,7 @@
 * Python: Added OBJECT REFCOUNT command ([#1485](https://github.com/aws/glide-for-redis/pull/1485))
 * Python: Added RENAMENX command ([#1492](https://github.com/aws/glide-for-redis/pull/1492))
 * Python: Added PFCOUNT command ([#1493](https://github.com/aws/glide-for-redis/pull/1493))
+* Python: Added PFMERGE command ([#1497](https://github.com/aws/glide-for-redis/pull/1497))
 
 ## 0.4.0 (2024-05-26)
 

python/python/glide/async_commands/core.py

@@ -3574,6 +3574,38 @@ async def pfcount(self, keys: List[str]) -> int:
             await self._execute_command(RequestType.PfCount, keys),
         )
 
+    async def pfmerge(self, destination: str, source_keys: List[str]) -> TOK:
+        """
+        Merges multiple HyperLogLog values into a unique value. If the destination variable exists, it is treated as one
+        of the source HyperLogLog data sets, otherwise a new HyperLogLog is created.
+
+        See https://valkey.io/commands/pfmerge for more details.
+
+        Note:
+            When in Cluster mode, all keys in `source_keys` and `destination` must map to the same hash slot.
+
+        Args:
+            destination (str): The key of the destination HyperLogLog where the merged data sets will be stored.
+            source_keys (List[str]): The keys of the HyperLogLog structures to be merged.
+
+        Returns:
+            OK: A simple OK response.
+
+        Examples:
+            >>> await client.pfadd("hll1", ["a", "b"])
+            >>> await client.pfadd("hll2", ["b", "c"])
+            >>> await client.pfmerge("new_hll", ["hll1", "hll2"])
+                OK  # The value of "hll1" merged with "hll2" was stored in "new_hll".
+            >>> await client.pfcount(["new_hll"])
+                3  # The approximated cardinality of "new_hll" is 3.
+        """
+        return cast(
+            TOK,
+            await self._execute_command(
+                RequestType.PfMerge, [destination] + source_keys
+            ),
+        )
+
     async def object_encoding(self, key: str) -> Optional[str]:
         """
         Returns the internal encoding for the Redis object stored at `key`.

python/python/glide/async_commands/transaction.py

@@ -2491,6 +2491,24 @@ def pfcount(self: TTransaction, keys: List[str]) -> TTransaction:
         """
         return self.append_command(RequestType.PfCount, keys)
 
+    def pfmerge(
+        self: TTransaction, destination: str, source_keys: List[str]
+    ) -> TTransaction:
+        """
+        Merges multiple HyperLogLog values into a unique value. If the destination variable exists, it is treated as one
+        of the source HyperLogLog data sets, otherwise a new HyperLogLog is created.
+
+        See https://valkey.io/commands/pfmerge for more details.
+
+        Args:
+            destination (str): The key of the destination HyperLogLog where the merged data sets will be stored.
+            source_keys (List[str]): The keys of the HyperLogLog structures to be merged.
+
+        Command response:
+            OK: A simple OK response.
+        """
+        return self.append_command(RequestType.PfMerge, [destination] + source_keys)
+
     def object_encoding(self: TTransaction, key: str) -> TTransaction:
         """
         Returns the internal encoding for the Redis object stored at `key`.

python/python/tests/test_async_client.py

@@ -3372,6 +3372,43 @@ async def test_pfcount(self, redis_client: TRedisClient):
         with pytest.raises(RequestError):
             await redis_client.pfcount([string_key])
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_pfmerge(self, redis_client: TRedisClient):
+        key1 = f"{{testKey}}:1-{get_random_string(10)}"
+        key2 = f"{{testKey}}:2-{get_random_string(10)}"
+        key3 = f"{{testKey}}:3-{get_random_string(10)}"
+        string_key = f"{{testKey}}:4-{get_random_string(10)}"
+        non_existing_key = f"{{testKey}}:5-{get_random_string(10)}"
+
+        assert await redis_client.pfadd(key1, ["a", "b", "c"]) == 1
+        assert await redis_client.pfadd(key2, ["b", "c", "d"]) == 1
+
+        # merge into new HyperLogLog data set
+        assert await redis_client.pfmerge(key3, [key1, key2]) == OK
+        assert await redis_client.pfcount([key3]) == 4
+
+        # merge into existing HyperLogLog data set
+        assert await redis_client.pfmerge(key1, [key2]) == OK
+        assert await redis_client.pfcount([key1]) == 4
+
+        # non-existing source key
+        assert await redis_client.pfmerge(key2, [key1, non_existing_key]) == OK
+        assert await redis_client.pfcount([key2]) == 4
+
+        # empty source key list
+        assert await redis_client.pfmerge(key1, []) == OK
+        assert await redis_client.pfcount([key1]) == 4
+
+        # source key exists, but it is not a HyperLogLog
+        assert await redis_client.set(string_key, "foo")
+        with pytest.raises(RequestError):
+            assert await redis_client.pfmerge(key3, [string_key])
+
+        # destination key exists, but it is not a HyperLogLog
+        with pytest.raises(RequestError):
+            assert await redis_client.pfmerge(string_key, [key3])
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_object_encoding(self, redis_client: TRedisClient):
@@ -3516,6 +3553,7 @@ async def test_multi_key_command_returns_cross_slot_error(
             redis_client.sdiffstore("abc", ["def", "ghi"]),
             redis_client.renamenx("abc", "def"),
             redis_client.pfcount(["def", "ghi"]),
+            redis_client.pfmerge("abc", ["def", "ghi"]),
         ]
 
         if not await check_if_server_version_lt(redis_client, "7.0.0"):

python/python/tests/test_transaction.py

@@ -304,6 +304,8 @@ async def transaction_test(
 
     transaction.pfadd(key10, ["a", "b", "c"])
     args.append(1)
+    transaction.pfmerge(key10, [])
+    args.append(OK)
     transaction.pfcount([key10])
     args.append(3)