Skip to content

Commit

Permalink
Python: add PFMERGE command (#322)
Browse files Browse the repository at this point in the history
  • Loading branch information
aaron-congo authored May 30, 2024
1 parent b2be117 commit 82f0c19
Show file tree
Hide file tree
Showing 5 changed files with 91 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 (TODO: add PR link)

## 0.4.0 (2024-05-26)

Expand Down
32 changes: 32 additions & 0 deletions python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down
18 changes: 18 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down
38 changes: 38 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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"):
Expand Down
2 changes: 2 additions & 0 deletions python/python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit 82f0c19

Please sign in to comment.