Skip to content

Commit fcc0c41

Browse files
authored
Support TDIGEST.MERGESTORE and make compression optional on TDIGEST.CREATE (#2319)
* support 2.4 * async test * skip tests * linters
1 parent 19cedab commit fcc0c41

File tree

3 files changed

+41
-1
lines changed

3 files changed

+41
-1
lines changed

redis/commands/bf/commands.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
TDIGEST_MIN = "TDIGEST.MIN"
5050
TDIGEST_MAX = "TDIGEST.MAX"
5151
TDIGEST_INFO = "TDIGEST.INFO"
52+
TDIGEST_MERGESTORE = "TDIGEST.MERGESTORE"
5253

5354

5455
class BFCommands:
@@ -344,7 +345,7 @@ def info(self, key):
344345

345346

346347
class TDigestCommands:
347-
def create(self, key, compression):
348+
def create(self, key, compression=100):
348349
"""
349350
Allocate the memory and initialize the t-digest.
350351
For more information see `TDIGEST.CREATE <https://redis.io/commands/tdigest.create>`_.
@@ -417,6 +418,19 @@ def info(self, key):
417418
""" # noqa
418419
return self.execute_command(TDIGEST_INFO, key)
419420

421+
def mergestore(self, dest_key, numkeys, *sourcekeys, compression=False):
422+
"""
423+
Merges all of the values from `sourcekeys` keys to `dest_key` sketch.
424+
If destination already exists, it is overwritten.
425+
426+
427+
For more information see `TDIGEST.MERGESTORE <https://redis.io/commands/tdigest.mergestore>`_.
428+
""" # noqa
429+
params = [dest_key, numkeys, *sourcekeys]
430+
if compression:
431+
params.extend(["COMPRESSION", compression])
432+
return self.execute_command(TDIGEST_MERGESTORE, *params)
433+
420434

421435
class CMSCommands:
422436
"""Count-Min Sketch Commands"""

tests/test_asyncio/test_bloom.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,19 @@ async def test_tdigest_cdf(modclient: redis.Redis):
393393
assert 0.9 == round(await modclient.tdigest().cdf("tDigest", 9.0), 1)
394394

395395

396+
@pytest.mark.redismod
397+
@pytest.mark.experimental
398+
@skip_ifmodversion_lt("2.4.0", "bf")
399+
async def test_tdigest_mergestore(modclient: redis.Redis):
400+
assert await modclient.tdigest().create("sourcekey1", 100)
401+
assert await modclient.tdigest().create("sourcekey2", 100)
402+
assert await modclient.tdigest().add("sourcekey1", [10], [1.0])
403+
assert await modclient.tdigest().add("sourcekey2", [50], [1.0])
404+
assert await modclient.tdigest().mergestore("dest", 2, "sourcekey1", "sourcekey2")
405+
assert await modclient.tdigest().max("dest") == 50
406+
assert await modclient.tdigest().min("dest") == 10
407+
408+
396409
# @pytest.mark.redismod
397410
# async def test_pipeline(modclient: redis.Redis):
398411
# pipeline = await modclient.bf().pipeline()

tests/test_bloom.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,19 @@ def test_tdigest_cdf(client):
388388
assert 0.9 == round(client.tdigest().cdf("tDigest", 9.0), 1)
389389

390390

391+
@pytest.mark.redismod
392+
@pytest.mark.experimental
393+
@skip_ifmodversion_lt("2.4.0", "bf")
394+
def test_tdigest_mergestore(client):
395+
assert client.tdigest().create("sourcekey1", 100)
396+
assert client.tdigest().create("sourcekey2", 100)
397+
assert client.tdigest().add("sourcekey1", [10], [1.0])
398+
assert client.tdigest().add("sourcekey2", [50], [1.0])
399+
assert client.tdigest().mergestore("destkey", 2, "sourcekey1", "sourcekey2")
400+
assert client.tdigest().max("destkey") == 50
401+
assert client.tdigest().min("destkey") == 10
402+
403+
391404
# @pytest.mark.redismod
392405
# def test_pipeline(client):
393406
# pipeline = client.bf().pipeline()

0 commit comments

Comments
 (0)