diff options
author | Avital Fine <79420960+AvitalFineRedis@users.noreply.github.com> | 2021-08-15 14:13:28 +0300 |
---|---|---|
committer | GitHub <noreply@github.com> | 2021-08-15 14:13:28 +0300 |
commit | 8c82e429224e0fb7f1739f741e63f23126e3c088 (patch) | |
tree | b7ff8f14d5725e05afdf5dafe4bea38704006caa | |
parent | 161774be431f2de3c44037698ac311735d559636 (diff) | |
download | redis-py-8c82e429224e0fb7f1739f741e63f23126e3c088.tar.gz |
Zunion (#1522)
* zinter
* change options in _zaggregate
* skip for previous versions
* add client function
* validate the aggregate value
* change options to get
* add more aggregate tests
* add weights guidance
-rwxr-xr-x | redis/client.py | 12 | ||||
-rw-r--r-- | tests/test_commands.py | 20 |
2 files changed, 31 insertions, 1 deletions
diff --git a/redis/client.py b/redis/client.py index 28e3ac1..f50c27e 100755 --- a/redis/client.py +++ b/redis/client.py @@ -615,7 +615,7 @@ class Redis: lambda r: r and set(r) or set() ), **string_keys_to_dict( - 'ZPOPMAX ZPOPMIN ZINTER ZDIFF ZRANGE ZRANGEBYSCORE ' + 'ZPOPMAX ZPOPMIN ZINTER ZDIFF ZUNION ZRANGE ZRANGEBYSCORE ' 'ZREVRANGE ZREVRANGEBYSCORE', zset_score_pairs ), **string_keys_to_dict('BZPOPMIN BZPOPMAX', \ @@ -3416,6 +3416,16 @@ class Redis: "Return the score of element ``value`` in sorted set ``name``" return self.execute_command('ZSCORE', name, value) + def zunion(self, keys, aggregate=None, withscores=False): + """ + Return the union of multiple sorted sets specified by ``keys``. + ``keys`` can be provided as dictionary of keys and their weights. + Scores will be aggregated based on the ``aggregate``, or SUM if + none is provided. + """ + return self._zaggregate('ZUNION', None, keys, aggregate, + withscores=withscores) + def zunionstore(self, dest, keys, aggregate=None): """ Union multiple sorted sets specified by ``keys`` into diff --git a/tests/test_commands.py b/tests/test_commands.py index 9a618aa..3c87dd9 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1813,6 +1813,26 @@ class TestRedisCommands: assert r.zscore('a', 'a2') == 2.0 assert r.zscore('a', 'a4') is None + @skip_if_server_version_lt('6.2.0') + def test_zunion(self, r): + r.zadd('a', {'a1': 1, 'a2': 1, 'a3': 1}) + r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) + r.zadd('c', {'a1': 6, 'a3': 5, 'a4': 4}) + # sum + assert r.zunion(['a', 'b', 'c']) == \ + [b'a2', b'a4', b'a3', b'a1'] + assert r.zunion(['a', 'b', 'c'], withscores=True) == \ + [(b'a2', 3), (b'a4', 4), (b'a3', 8), (b'a1', 9)] + # max + assert r.zunion(['a', 'b', 'c'], aggregate='MAX', withscores=True)\ + == [(b'a2', 2), (b'a4', 4), (b'a3', 5), (b'a1', 6)] + # min + assert r.zunion(['a', 'b', 'c'], aggregate='MIN', withscores=True)\ + == [(b'a1', 1), (b'a2', 1), (b'a3', 1), (b'a4', 4)] + # with weight + assert r.zunion({'a': 1, 'b': 2, 'c': 3}, withscores=True)\ + == [(b'a2', 5), (b'a4', 12), (b'a3', 20), (b'a1', 23)] + def test_zunionstore_sum(self, r): r.zadd('a', {'a1': 1, 'a2': 1, 'a3': 1}) r.zadd('b', {'a1': 2, 'a2': 2, 'a3': 2}) |