ZUNIONRANGESCORE

test.py

import redis
r = redis.StrictRedis()

r.zadd('k1', 1, 'a', 2, 'b', 3, 'c')
r.zadd('k2', 2, 'a', 3, 'b', 4, 'c')
r.zadd('k3', 3, 'a', 4, 'b', 5, 'c')

from zunion_range_score import zunion_range_score as zunionrangescore

range_scores = { 'k1':(1, 5), 'k2':(1, 5), 'k3':(1, 5) }
print zunionrangescore(r, range_scores)

range_scores = { 'k1':(1, 5), 'k2':(1, 5), 'k3':(1, 4) }
print zunionrangescore(r, range_scores)

range_scores = { 'k1':(2, 5), 'k2':(3, 5), 'k3':(4, 5) }
print zunionrangescore(r, range_scores)

range_scores = { 'k1':(4, 5), 'k2':(3, 5), 'k3':(2, 5) }
print zunionrangescore(r, range_scores)

zunion_range_score.py

'''
Written on March 28, 2014 by Josiah Carlson
Released into the public domain

ZUNIONRANGESCORE:
Zunion Range Score performs a Redis ZUNIONSTORE operation, selecting *only
those items in the provided ranges. The scores are added. Proof of concept.

Warning: untested, use at your own risk.

'''

import json
from uuid import uuid4

def zunion_range_score(conn, ranges_scores):
    '''
    Args:
        conn - Redis connection
        ranges_scores - a dictionary mapping "key" to (min, max) ranges of
                        scores (passed directly to Redis)

    This function returns the union of all keys with the ranges of scores to
    pull from. Intersection and mixed-mode left as an exersize to the reader.
    '''
    tempkey = str(uuid4())
    keys, argv = zip(*sorted(ranges_scores.items()))
    keys = list(keys) + [tempkey]
    argv = [json.dumps(argv)]
    zunionrangescore = conn.register_script(_union_range_score_lua)
    print keys, argv
    return zunionrangescore(keys, argv)
    
_union_range_score_lua = '''
-- KEYS: [key1, key2, ..., keyn, tempkey]
-- ARGV: [JSON([[min1, max1], [min2, max2], ...])]

local _ARGV = cjson.decode(ARGV[1])
local dest = KEYS[#KEYS]

for i = 1, #_ARGV do
    local _a = _ARGV[i]; table.insert(_a, 'withscores')
    local chunk = redis.call('ZRANGEBYSCORE', KEYS[i], unpack(_a))
    if #chunk > 0 then
        for j = 1, #chunk, 2 do
            redis.call('ZINCRBY', dest, chunk[j+1], chunk[j])
        end
    end
end

local ret = redis.call('ZRANGE', dest, 0, -1, 'withscores')
redis.call('DEL', dest)
return ret
'''