From 76da3d4df38457ca030b9ed2b9f4ada67ebe1b50 Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Sat, 24 Jul 2021 13:31:30 +0300 Subject: [PATCH 1/8] zinter --- redis/client.py | 38 ++++++++++++++++++++++++++++++-------- tests/test_commands.py | 18 ++++++++++++++++++ 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/redis/client.py b/redis/client.py index 57932b9f73..19d0d2cc85 100755 --- a/redis/client.py +++ b/redis/client.py @@ -595,8 +595,8 @@ class Redis: lambda r: r and set(r) or set() ), **string_keys_to_dict( - 'ZPOPMAX ZPOPMIN ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE', - zset_score_pairs + 'ZPOPMAX ZPOPMIN ZINTER ZRANGE ZRANGEBYSCORE ZREVRANGE ' + 'ZREVRANGEBYSCORE', zset_score_pairs ), **string_keys_to_dict('BZPOPMIN BZPOPMAX', \ lambda r: @@ -2906,11 +2906,27 @@ def zincrby(self, name, amount, value): "Increment the score of ``value`` in sorted set ``name`` by ``amount``" return self.execute_command('ZINCRBY', name, amount, value) + def zinter(self, keys, aggregate=None, withscores=False): + """ + Return the intersect of multiple sorted sets specified by ``keys``. + With the ``aggregate`` option, it is possible to specify how the + results of the union are aggregated. This option defaults to SUM, + where the score of an element is summed across the inputs where it + exists. When this option is set to either MIN or MAX, the resulting + set will contain the minimum or maximum score of an element across + the inputs where it exists. + """ + return self._zaggregate('ZINTER', None, keys, aggregate, withscores) + def zinterstore(self, dest, keys, aggregate=None): """ - Intersect multiple sorted sets specified by ``keys`` into - a new sorted set, ``dest``. Scores in the destination will be - aggregated based on the ``aggregate``, or SUM if none is provided. + Intersect multiple sorted sets specified by ``keys`` into a new + sorted set, ``dest``. Scores in the destination will be aggregated + based on the ``aggregate``. This option defaults to SUM, where the + score of an element is summed across the inputs where it exists. + When this option is set to either MIN or MAX, the resulting set will + contain the minimum or maximum score of an element across the inputs + where it exists. """ return self._zaggregate('ZINTERSTORE', dest, keys, aggregate) @@ -3169,8 +3185,12 @@ def zunionstore(self, dest, keys, aggregate=None): """ return self._zaggregate('ZUNIONSTORE', dest, keys, aggregate) - def _zaggregate(self, command, dest, keys, aggregate=None): - pieces = [command, dest, len(keys)] + def _zaggregate(self, command, dest, keys, aggregate=None, + withscores=False): + pieces = [command] + if dest is not None: + pieces.append(dest) + pieces.append(len(keys)) if isinstance(keys, dict): keys, weights = keys.keys(), keys.values() else: @@ -3182,7 +3202,9 @@ def _zaggregate(self, command, dest, keys, aggregate=None): if aggregate: pieces.append(b'AGGREGATE') pieces.append(aggregate) - return self.execute_command(*pieces) + if withscores: + pieces.append(b'WITHSCORES') + return self.execute_command(*pieces, withscores=withscores) # HYPERLOGLOG COMMANDS def pfadd(self, name, *values): diff --git a/tests/test_commands.py b/tests/test_commands.py index 27117e3188..615d5de42c 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1456,6 +1456,24 @@ def test_zlexcount(self, r): assert r.zlexcount('a', '-', '+') == 7 assert r.zlexcount('a', '[b', '[f') == 5 + def test_zinter(self, r): + r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 1}) + r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) + r.zadd('c', {'a1': 6, 'a3': 5, 'a4': 4}) + assert r.zinter(['a', 'b', 'c']) == [b'a3', b'a1'] + # aggregate with SUM + assert r.zinter(['a', 'b', 'c'], withscores=True) \ + == [(b'a3', 8), (b'a1', 9)] + # aggregate with MAX + assert r.zinter(['a', 'b', 'c'], aggregate='MAX', withscores=True) \ + == [(b'a3', 5), (b'a1', 6)] + # aggregate with MIN + assert r.zinter(['a', 'b', 'c'], aggregate='MIN', withscores=True) \ + == [(b'a1', 1), (b'a3', 1)] + # with weights + assert r.zinter({'a': 1, 'b': 2, 'c': 3}, withscores=True) \ + == [(b'a3', 20), (b'a1', 23)] + def test_zinterstore_sum(self, r): r.zadd('a', {'a1': 1, 'a2': 1, 'a3': 1}) r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) From f7c5fb9054edfebbcf192eccc947c3f2ba4a147a Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Sun, 25 Jul 2021 09:20:22 +0300 Subject: [PATCH 2/8] change options in _zaggregate --- redis/client.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/redis/client.py b/redis/client.py index 19d0d2cc85..7f3a9bbe00 100755 --- a/redis/client.py +++ b/redis/client.py @@ -2916,7 +2916,7 @@ def zinter(self, keys, aggregate=None, withscores=False): set will contain the minimum or maximum score of an element across the inputs where it exists. """ - return self._zaggregate('ZINTER', None, keys, aggregate, withscores) + return self._zaggregate('ZINTER', None, keys, aggregate, withscores=withscores) def zinterstore(self, dest, keys, aggregate=None): """ @@ -3186,7 +3186,7 @@ def zunionstore(self, dest, keys, aggregate=None): return self._zaggregate('ZUNIONSTORE', dest, keys, aggregate) def _zaggregate(self, command, dest, keys, aggregate=None, - withscores=False): + **options): pieces = [command] if dest is not None: pieces.append(dest) @@ -3202,9 +3202,9 @@ def _zaggregate(self, command, dest, keys, aggregate=None, if aggregate: pieces.append(b'AGGREGATE') pieces.append(aggregate) - if withscores: + if 'withscores' in options.keys() and options['withscores']: pieces.append(b'WITHSCORES') - return self.execute_command(*pieces, withscores=withscores) + return self.execute_command(*pieces, **options) # HYPERLOGLOG COMMANDS def pfadd(self, name, *values): From 64f73f8a69106e5b0ad7c9d5319a5593461ab160 Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Sun, 25 Jul 2021 10:43:20 +0300 Subject: [PATCH 3/8] skip for previous versions --- tests/test_commands.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_commands.py b/tests/test_commands.py index 615d5de42c..2a8f3caea4 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1456,6 +1456,7 @@ def test_zlexcount(self, r): assert r.zlexcount('a', '-', '+') == 7 assert r.zlexcount('a', '[b', '[f') == 5 + @skip_if_server_version_lt('6.2.0') def test_zinter(self, r): r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 1}) r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) From 983d4c92da728029cfc5d382ef377c58189bb5f6 Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Sun, 25 Jul 2021 10:54:19 +0300 Subject: [PATCH 4/8] flake8 --- redis/client.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/redis/client.py b/redis/client.py index 7f3a9bbe00..1ee426982a 100755 --- a/redis/client.py +++ b/redis/client.py @@ -2916,7 +2916,8 @@ def zinter(self, keys, aggregate=None, withscores=False): set will contain the minimum or maximum score of an element across the inputs where it exists. """ - return self._zaggregate('ZINTER', None, keys, aggregate, withscores=withscores) + return self._zaggregate('ZINTER', None, keys, aggregate, + withscores=withscores) def zinterstore(self, dest, keys, aggregate=None): """ From eb9c6161019df079a95d2a217988f2dadcf1f277 Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Mon, 26 Jul 2021 13:50:31 +0300 Subject: [PATCH 5/8] validate the aggregate value --- redis/client.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/redis/client.py b/redis/client.py index 1ee426982a..bba70e242b 100755 --- a/redis/client.py +++ b/redis/client.py @@ -3201,8 +3201,11 @@ def _zaggregate(self, command, dest, keys, aggregate=None, pieces.append(b'WEIGHTS') pieces.extend(weights) if aggregate: - pieces.append(b'AGGREGATE') - pieces.append(aggregate) + if aggregate.upper() in ['SUM', 'MIN', 'MAX']: + pieces.append(b'AGGREGATE') + pieces.append(aggregate) + else: + raise DataError("aggregate can be sum, min or max.") if 'withscores' in options.keys() and options['withscores']: pieces.append(b'WITHSCORES') return self.execute_command(*pieces, **options) From 9aee43ad336f64c4d0470f69c6ade882f17f130d Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Mon, 26 Jul 2021 14:08:32 +0300 Subject: [PATCH 6/8] invalid aggregation --- tests/test_commands.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_commands.py b/tests/test_commands.py index 2a8f3caea4..84cd5a1231 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1462,6 +1462,9 @@ def test_zinter(self, r): r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) r.zadd('c', {'a1': 6, 'a3': 5, 'a4': 4}) assert r.zinter(['a', 'b', 'c']) == [b'a3', b'a1'] + # invalid aggregation + with pytest.raises(exceptions.DataError) as e: + r.zinter(['a', 'b', 'c'], aggregate='foo', withscores=True) # aggregate with SUM assert r.zinter(['a', 'b', 'c'], withscores=True) \ == [(b'a3', 8), (b'a1', 9)] From f1d36f56384f1360c9a925598bf80f8b8f3e7277 Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Mon, 26 Jul 2021 14:08:32 +0300 Subject: [PATCH 7/8] invalid aggregation --- tests/test_commands.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_commands.py b/tests/test_commands.py index 2a8f3caea4..ee5bd4eb74 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1462,6 +1462,9 @@ def test_zinter(self, r): r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) r.zadd('c', {'a1': 6, 'a3': 5, 'a4': 4}) assert r.zinter(['a', 'b', 'c']) == [b'a3', b'a1'] + # invalid aggregation + with pytest.raises(exceptions.DataError): + r.zinter(['a', 'b', 'c'], aggregate='foo', withscores=True) # aggregate with SUM assert r.zinter(['a', 'b', 'c'], withscores=True) \ == [(b'a3', 8), (b'a1', 9)] From bce93654d48b8c452e2ac496d8b0f4a3c6bd05af Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Wed, 28 Jul 2021 17:03:18 +0300 Subject: [PATCH 8/8] change options to get --- redis/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redis/client.py b/redis/client.py index bba70e242b..f2bb39e992 100755 --- a/redis/client.py +++ b/redis/client.py @@ -3206,7 +3206,7 @@ def _zaggregate(self, command, dest, keys, aggregate=None, pieces.append(aggregate) else: raise DataError("aggregate can be sum, min or max.") - if 'withscores' in options.keys() and options['withscores']: + if options.get('withscores', False): pieces.append(b'WITHSCORES') return self.execute_command(*pieces, **options)