Skip to content
Browse files

added support for zinter and zunion

  • Loading branch information...
1 parent 60063e8 commit f4c8393113e2205eb8eeddeed10d42ad70d648ef @andymccurdy andymccurdy committed Apr 28, 2010
Showing with 86 additions and 0 deletions.
  1. +32 −0 redis/client.py
  2. +54 −0 tests/server_commands.py
View
32 redis/client.py
@@ -896,6 +896,14 @@ def zincrby(self, name, value, amount=1):
"Increment the score of ``value`` in sorted set ``name`` by ``amount``"
return self.execute_command('ZINCRBY', name, amount, value)
+ def zinter(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.
+ """
+ return self._zaggregate('ZINTER', dest, keys, aggregate)
+
def zrange(self, name, start, end, desc=False, withscores=False):
"""
Return a range of values from sorted set ``name`` between
@@ -980,6 +988,30 @@ def zscore(self, name, value):
"Return the score of element ``value`` in sorted set ``name``"
return self.execute_command('ZSCORE', name, value)
+ def zunion(self, dest, keys, aggregate=None):
+ """
+ Union 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.
+ """
+ return self._zaggregate('ZUNION', dest, keys, aggregate)
+
+ def _zaggregate(self, command, dest, keys, aggregate=None):
+ pieces = [command, dest, len(keys)]
+ if isinstance(keys, dict):
+ items = keys.items()
+ keys = [i[0] for i in items]
+ weights = [i[1] for i in items]
+ else:
+ weights = None
+ pieces.extend(keys)
+ if weights:
+ pieces.append('WEIGHTS')
+ pieces.extend(weights)
+ if aggregate:
+ pieces.append('AGGREGATE')
+ pieces.append(aggregate)
+ return self.execute_command(*pieces)
#### HASH COMMANDS ####
def hdel(self, name, key):
View
54 tests/server_commands.py
@@ -611,6 +611,33 @@ def test_zincrby(self):
self.assertEquals(self.client.zscore('a', 'a2'), 3.0)
self.assertEquals(self.client.zscore('a', 'a3'), 8.0)
+ def test_zinter(self):
+ self.make_zset('a', {'a1': 1, 'a2': 1, 'a3': 1})
+ self.make_zset('b', {'a1': 2, 'a3': 2, 'a4': 2})
+ self.make_zset('c', {'a1': 6, 'a3': 5, 'a4': 4})
+
+ # sum, no weight
+ self.assert_(self.client.zinter('z', ['a', 'b', 'c']))
+ self.assertEquals(
+ self.client.zrange('z', 0, -1, withscores=True),
+ [('a3', 8), ('a1', 9)]
+ )
+
+ # max, no weight
+ self.assert_(self.client.zinter('z', ['a', 'b', 'c'], aggregate='MAX'))
+ self.assertEquals(
+ self.client.zrange('z', 0, -1, withscores=True),
+ [('a3', 5), ('a1', 6)]
+ )
+
+ # with weight
+ self.assert_(self.client.zinter('z', {'a': 1, 'b': 2, 'c': 3}))
+ self.assertEquals(
+ self.client.zrange('z', 0, -1, withscores=True),
+ [('a3', 20), ('a1', 23)]
+ )
+
+
def test_zrange(self):
# key is not a zset
self.client['a'] = 'a'
@@ -724,6 +751,33 @@ def test_zscore(self):
# test a non-existant member
self.assertEquals(self.client.zscore('a', 'a4'), None)
+ def test_zunion(self):
+ self.make_zset('a', {'a1': 1, 'a2': 1, 'a3': 1})
+ self.make_zset('b', {'a1': 2, 'a3': 2, 'a4': 2})
+ self.make_zset('c', {'a1': 6, 'a4': 5, 'a5': 4})
+
+ # sum, no weight
+ self.assert_(self.client.zunion('z', ['a', 'b', 'c']))
+ self.assertEquals(
+ self.client.zrange('z', 0, -1, withscores=True),
+ [('a2', 1), ('a3', 3), ('a5', 4), ('a4', 7), ('a1', 9)]
+ )
+
+ # max, no weight
+ self.assert_(self.client.zunion('z', ['a', 'b', 'c'], aggregate='MAX'))
+ self.assertEquals(
+ self.client.zrange('z', 0, -1, withscores=True),
+ [('a2', 1), ('a3', 2), ('a5', 4), ('a4', 5), ('a1', 6)]
+ )
+
+ # with weight
+ self.assert_(self.client.zunion('z', {'a': 1, 'b': 2, 'c': 3}))
+ self.assertEquals(
+ self.client.zrange('z', 0, -1, withscores=True),
+ [('a2', 1), ('a3', 5), ('a5', 12), ('a4', 19), ('a1', 23)]
+ )
+
+
# HASHES
def make_hash(self, key, d):
for k,v in d.iteritems():

0 comments on commit f4c8393

Please sign in to comment.
Something went wrong with that request. Please try again.