diff options
-rw-r--r-- | redis/commands/bf/__init__.py | 3 | ||||
-rw-r--r-- | redis/commands/bf/commands.py | 9 | ||||
-rw-r--r-- | redis/commands/bf/utils.py | 3 | ||||
-rw-r--r-- | tests/test_asyncio/test_bloom.py | 28 | ||||
-rw-r--r-- | tests/test_bloom.py | 20 |
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 |