Add support for TDIGEST.QUANTILE extensions (#2317)

* Add support for TDIGEST.QUANTILE extensions

* linters

* linters & utils

* Update test_bloom.py

redis/commands/bf/__init__.py

@@ -3,6 +3,7 @@
 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 @@ def __init__(self, client, **kwargs):
             # 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,

redis/commands/bf/commands.py

@@ -393,13 +393,14 @@ def max(self, key):
         """  # 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):
         """

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]

tests/test_asyncio/test_bloom.py

@@ -3,6 +3,7 @@
 import redis.asyncio as redis
 from redis.exceptions import ModuleError, RedisError
 from redis.utils import HIREDIS_AVAILABLE
+from tests.conftest import skip_ifmodversion_lt
 
 pytestmark = pytest.mark.asyncio
 
@@ -357,22 +358,31 @@ 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
     assert await modclient.tdigest().add(
         "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

tests/test_bloom.py

@@ -4,6 +4,8 @@
 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,18 +356,26 @@ 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
     assert client.tdigest().add(
         "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