Instantly share code, notes, and snippets.

Embed
What would you like to do?
A method to get sql-like multiple-column order by in Redis
'''
sort_zset_cols.py
Copyright 2013 Josiah Carlson
Released into the public domain.
'''
'''
Let's imagine that there are 3 restaurants with price, score, distance info
being:
(50, 3, 20)
(40, 4, 10)
(40, 5, 50)
With the default sort, this method will create aggregate scores for the
above values as follows:
20507
10406
50405
If you were to sort that information, you would get:
10406, 20507, 50405
Which is the order you want: lowest distance, lowest price, highest score.
Notes:
* data_key is a key that currently stores the result of any
pre-filtering that has been done on the list of restaurants already, and
will be overwritten
* You may want to edit the keys that are provided as the 'sort' field, as
sortable data is not likely to be in keys with those names.
'''
import math
import warnings
def sort_zset_cols(conn, result_key, sort=('dist', 'price', '-score')):
current_multiplier = 1
keys = {result_key: 0}
sort = list(reversed(sort))
# Gets the max/min values in a sort column
pipe = conn.pipeline(True)
for sort_col in sort:
pipe.zrange(sort_col, 0, 0, withscores=True)
pipe.zrange(sort_col, -1, -1, withscores=True)
ranges = pipe.execute()
for i, sort_col in enumerate(sort):
# Auto-scaling for negative values
low, high = ranges[i*2][1], ranges[i*2+1][1]
maxv = int(math.ceil(max(abs(low), abs(high))))
# Adjusts the weights based on the magnitude and sort order of the
# column
old_multiplier = current_multiplier
desc = sort_col.startswith('-')
sort_col = sort_col.lstrip('-')
current_multiplier *= maxv
# Assign the sort key a weight based on all of the lower-priority
# sort columns
keys[sort_col] = -old_multiplier if desc else old_multiplier
if current_multiplier >= 2**53:
warnings.warn("The total range of your values is outside the "
"available score precision, and your sort may not be precise")
# The sort results are available in the passed result_key
return conn.zinterstore(result_key, keys)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment