Support TDIGEST.MERGESTORE and make compression optional on TDIGEST.CREATE (#2319)

* support 2.4

* async test

* skip tests

* linters

3 files changed, 41 insertions, 1 deletions

diff --git a/redis/commands/bf/commands.py b/redis/commands/bf/commands.py
index 7d36b93..84a6b5f 100644
--- a/redis/commands/bf/commands.py
+++ b/redis/commands/bf/commands.py
@@ -49,6 +49,7 @@ TDIGEST_QUANTILE = "TDIGEST.QUANTILE"
 TDIGEST_MIN = "TDIGEST.MIN"
 TDIGEST_MAX = "TDIGEST.MAX"
 TDIGEST_INFO = "TDIGEST.INFO"
+TDIGEST_MERGESTORE = "TDIGEST.MERGESTORE"
 
 
 class BFCommands:
@@ -344,7 +345,7 @@ class TOPKCommands:
 
 
 class TDigestCommands:
-    def create(self, key, compression):
+    def create(self, key, compression=100):
         """
         Allocate the memory and initialize the t-digest.
         For more information see `TDIGEST.CREATE <https://redis.io/commands/tdigest.create>`_.
@@ -417,6 +418,19 @@ class TDigestCommands:
         """  # noqa
         return self.execute_command(TDIGEST_INFO, key)
 
+    def mergestore(self, dest_key, numkeys, *sourcekeys, compression=False):
+        """
+        Merges all of the values from `sourcekeys` keys to `dest_key` sketch.
+        If destination already exists, it is overwritten.
+
+
+        For more information see `TDIGEST.MERGESTORE <https://redis.io/commands/tdigest.mergestore>`_.
+        """  # noqa
+        params = [dest_key, numkeys, *sourcekeys]
+        if compression:
+            params.extend(["COMPRESSION", compression])
+        return self.execute_command(TDIGEST_MERGESTORE, *params)
+
 
 class CMSCommands:
     """Count-Min Sketch Commands"""
diff --git a/tests/test_asyncio/test_bloom.py b/tests/test_asyncio/test_bloom.py
index fe3bd76..15b22b3 100644
--- a/tests/test_asyncio/test_bloom.py
+++ b/tests/test_asyncio/test_bloom.py
@@ -393,6 +393,19 @@ async def test_tdigest_cdf(modclient: redis.Redis):
     assert 0.9 == round(await modclient.tdigest().cdf("tDigest", 9.0), 1)
 
 
+@pytest.mark.redismod
+@pytest.mark.experimental
+@skip_ifmodversion_lt("2.4.0", "bf")
+async def test_tdigest_mergestore(modclient: redis.Redis):
+    assert await modclient.tdigest().create("sourcekey1", 100)
+    assert await modclient.tdigest().create("sourcekey2", 100)
+    assert await modclient.tdigest().add("sourcekey1", [10], [1.0])
+    assert await modclient.tdigest().add("sourcekey2", [50], [1.0])
+    assert await modclient.tdigest().mergestore("dest", 2, "sourcekey1", "sourcekey2")
+    assert await modclient.tdigest().max("dest") == 50
+    assert await modclient.tdigest().min("dest") == 10
+
+
 # @pytest.mark.redismod
 # async def test_pipeline(modclient: redis.Redis):
 #     pipeline = await modclient.bf().pipeline()
diff --git a/tests/test_bloom.py b/tests/test_bloom.py
index 340e48e..1191542 100644
--- a/tests/test_bloom.py
+++ b/tests/test_bloom.py
@@ -388,6 +388,19 @@ def test_tdigest_cdf(client):
     assert 0.9 == round(client.tdigest().cdf("tDigest", 9.0), 1)
 
 
+@pytest.mark.redismod
+@pytest.mark.experimental
+@skip_ifmodversion_lt("2.4.0", "bf")
+def test_tdigest_mergestore(client):
+    assert client.tdigest().create("sourcekey1", 100)
+    assert client.tdigest().create("sourcekey2", 100)
+    assert client.tdigest().add("sourcekey1", [10], [1.0])
+    assert client.tdigest().add("sourcekey2", [50], [1.0])
+    assert client.tdigest().mergestore("destkey", 2, "sourcekey1", "sourcekey2")
+    assert client.tdigest().max("destkey") == 50
+    assert client.tdigest().min("destkey") == 10
+
+
 # @pytest.mark.redismod
 # def test_pipeline(client):
 #     pipeline = client.bf().pipeline()