Skip to content

Instantly share code, notes, and snippets.

@josiahcarlson
Last active August 29, 2015 13:57
Show Gist options
  • Save josiahcarlson/9849543 to your computer and use it in GitHub Desktop.
Save josiahcarlson/9849543 to your computer and use it in GitHub Desktop.
Zunion Range Score, an interesting scripting solution.
'''
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.
'''
from hashlib import sha1
import json
from uuid import uuid4
def zunion_range_score(conn, ranges_scores, sha=[]):
'''
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()))
argv = [json.dumps(argv)]
if sha:
return conn.evalsha(sha[0], keys=keys + [tempkey], argv=argv)
sha.append(sha1(_union_range_score_lua).hexdigest())
return conn.eval(_union_range_score_lua, keys=keys + [tempkey], argv=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.remove()
for i = 1, #_ARGV do
local _a = _ARGV[i]; table.insert(_a, 'withscores')
local chunk = redis.call('ZRANGEBYSCORE', KEYS[i], unpack(a))
for i = 1, #chunk, 2 do
-- Hah, the inversion of ordering for scores and members in ZSET
-- reading vs. writing bites us in the butt.
chunk[i], chunk[i+1] = chunk[i+1], chunk[i]
end
redis.call('ZADD', unpack(chunk))
end
local ret = redis.call('ZRANGE', dest, 0, -1, withscores)
redis.call('DEL', dest)
return ret
'''
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment