Add support for `TDIGEST.QUANTILE` extensions (#2317)

* Add support for TDIGEST.QUANTILE extensions

* linters

* linters & utils

* Update test_bloom.py

5 files changed, 44 insertions, 19 deletions

diff --git a/redis/commands/bf/__init__.py b/redis/commands/bf/__init__.py
index d62d8a0..3169ba2 100644
--- a/redis/commands/bf/__init__.py
+++ b/redis/commands/bf/__init__.py
@@ -3,6 +3,7 @@ from redis.client import bool_ok
 from ..helpers import parse_to_list
 from .commands import *  # noqa
 from .info import BFInfo, CFInfo, CMSInfo, TDigestInfo, TopKInfo
+from .utils import parse_tdigest_quantile
 
 
 class AbstractBloom(object):
@@ -166,7 +167,7 @@ class TDigestBloom(TDigestCommands, AbstractBloom):
             # TDIGEST_ADD: spaceHolder,
             # TDIGEST_MERGE: spaceHolder,
             TDIGEST_CDF: float,
-            TDIGEST_QUANTILE: float,
+            TDIGEST_QUANTILE: parse_tdigest_quantile,
             TDIGEST_MIN: float,
             TDIGEST_MAX: float,
             TDIGEST_INFO: TDigestInfo,
diff --git a/redis/commands/bf/commands.py b/redis/commands/bf/commands.py
index baf0130..7d36b93 100644
--- a/redis/commands/bf/commands.py
+++ b/redis/commands/bf/commands.py
@@ -393,13 +393,14 @@ class TDigestCommands:
         """  # noqa
         return self.execute_command(TDIGEST_MAX, key)
 
-    def quantile(self, key, quantile):
+    def quantile(self, key, quantile, *quantiles):
         """
-        Return double value estimate of the cutoff such that a specified fraction of the data
-        added to this TDigest would be less than or equal to the cutoff.
+        Returns estimates of one or more cutoffs such that a specified fraction of the
+        observations added to this t-digest would be less than or equal to each of the
+        specified cutoffs. (Multiple quantiles can be returned with one call)
         For more information see `TDIGEST.QUANTILE <https://redis.io/commands/tdigest.quantile>`_.
         """  # noqa
-        return self.execute_command(TDIGEST_QUANTILE, key, quantile)
+        return self.execute_command(TDIGEST_QUANTILE, key, quantile, *quantiles)
 
     def cdf(self, key, value):
         """
diff --git a/redis/commands/bf/utils.py b/redis/commands/bf/utils.py
new file mode 100644
index 0000000..21dcfa7
--- /dev/null
+++ b/redis/commands/bf/utils.py
@@ -0,0 +1,3 @@
+def parse_tdigest_quantile(response):
+    """Parse TDIGEST.QUANTILE response."""
+    return [float(x) for x in response]
diff --git a/tests/test_asyncio/test_bloom.py b/tests/test_asyncio/test_bloom.py
index 2c3617f..fe3bd76 100644
--- a/tests/test_asyncio/test_bloom.py
+++ b/tests/test_asyncio/test_bloom.py
@@ -3,6 +3,7 @@ import pytest
 import redis.asyncio as redis
 from redis.exceptions import ModuleError, RedisError
 from redis.utils import HIREDIS_AVAILABLE
+from tests.conftest import skip_ifmodversion_lt
 
 
 def intlist(obj):
@@ -355,6 +356,7 @@ async def test_tdigest_min_and_max(modclient: redis.Redis):
 
 @pytest.mark.redismod
 @pytest.mark.experimental
+@skip_ifmodversion_lt("2.4.0", "bf")
 async def test_tdigest_quantile(modclient: redis.Redis):
     assert await modclient.tdigest().create("tDigest", 500)
     # insert data-points into sketch
@@ -362,15 +364,23 @@ async def test_tdigest_quantile(modclient: redis.Redis):
         "tDigest", list([x * 0.01 for x in range(1, 10000)]), [1.0] * 10000
     )
     # assert min min/max have same result as quantile 0 and 1
-    assert await modclient.tdigest().max(
-        "tDigest"
-    ) == await modclient.tdigest().quantile("tDigest", 1.0)
-    assert await modclient.tdigest().min(
-        "tDigest"
-    ) == await modclient.tdigest().quantile("tDigest", 0.0)
-
-    assert 1.0 == round(await modclient.tdigest().quantile("tDigest", 0.01), 2)
-    assert 99.0 == round(await modclient.tdigest().quantile("tDigest", 0.99), 2)
+    assert (
+        await modclient.tdigest().max("tDigest")
+        == (await modclient.tdigest().quantile("tDigest", 1.0))[1]
+    )
+    assert (
+        await modclient.tdigest().min("tDigest")
+        == (await modclient.tdigest().quantile("tDigest", 0.0))[1]
+    )
+
+    assert 1.0 == round((await modclient.tdigest().quantile("tDigest", 0.01))[1], 2)
+    assert 99.0 == round((await modclient.tdigest().quantile("tDigest", 0.99))[1], 2)
+
+    # test multiple quantiles
+    assert await modclient.tdigest().create("t-digest", 100)
+    assert await modclient.tdigest().add("t-digest", [1, 2, 3, 4, 5], [1.0] * 5)
+    res = await modclient.tdigest().quantile("t-digest", 0.5, 0.8)
+    assert [0.5, 3.0, 0.8, 5.0] == res
 
 
 @pytest.mark.redismod
diff --git a/tests/test_bloom.py b/tests/test_bloom.py
index 1f8201c..340e48e 100644
--- a/tests/test_bloom.py
+++ b/tests/test_bloom.py
@@ -4,6 +4,8 @@ import redis.commands.bf
 from redis.exceptions import ModuleError, RedisError
 from redis.utils import HIREDIS_AVAILABLE
 
+from .conftest import skip_ifmodversion_lt
+
 
 def intlist(obj):
     return [int(v) for v in obj]
@@ -354,6 +356,7 @@ def test_tdigest_min_and_max(client):
 
 @pytest.mark.redismod
 @pytest.mark.experimental
+@skip_ifmodversion_lt("2.4.0", "bf")
 def test_tdigest_quantile(client):
     assert client.tdigest().create("tDigest", 500)
     # insert data-points into sketch
@@ -361,11 +364,18 @@ def test_tdigest_quantile(client):
         "tDigest", list([x * 0.01 for x in range(1, 10000)]), [1.0] * 10000
     )
     # assert min min/max have same result as quantile 0 and 1
-    assert client.tdigest().max("tDigest") == client.tdigest().quantile("tDigest", 1.0)
-    assert client.tdigest().min("tDigest") == client.tdigest().quantile("tDigest", 0.0)
-
-    assert 1.0 == round(client.tdigest().quantile("tDigest", 0.01), 2)
-    assert 99.0 == round(client.tdigest().quantile("tDigest", 0.99), 2)
+    res = client.tdigest().quantile("tDigest", 1.0)
+    assert client.tdigest().max("tDigest") == res[1]
+    res = client.tdigest().quantile("tDigest", 0.0)
+    assert client.tdigest().min("tDigest") == res[1]
+
+    assert 1.0 == round(client.tdigest().quantile("tDigest", 0.01)[1], 2)
+    assert 99.0 == round(client.tdigest().quantile("tDigest", 0.99)[1], 2)
+
+    # test multiple quantiles
+    assert client.tdigest().create("t-digest", 100)
+    assert client.tdigest().add("t-digest", [1, 2, 3, 4, 5], [1.0] * 5)
+    assert [0.5, 3.0, 0.8, 5.0] == client.tdigest().quantile("t-digest", 0.5, 0.8)
 
 
 @pytest.mark.redismod