"SORT by nosort" (skip sorting) respect sorted set ordering.

When SORT is called with the option BY set to a string constant not
inclduing the wildcard character "*", there is no way to sort the output
so any ordering is valid. This allows the SORT internals to optimize its
work and don't really sort the output at all.

However it was odd that this option was not able to retain the natural
order of a sorted set. This feature was requested by users multiple
times as sometimes to call SORT with GET against sorted sets as a way to
mass-fetch objects can be handy.

This commit introduces two things:

1) The ability of SORT to return sorted sets elements in their natural
ordering when `BY nosort` is specified, accordingly to `DESC / ASC` options.
2) The ability of SORT to optimize this case further if LIMIT is passed
as well, avoiding to really fetch the whole sorted set, but directly
obtaining the specified range.

Because in this case the sorting is always deterministic, no
post-sorting activity is performed when SORT is called from a Lua
script.

This commit fixes issue #98.

2 files changed, 129 insertions, 20 deletions

diff --git a/src/sort.c b/src/sort.c
index e5178cd0d..d18a52959 100644
--- a/src/sort.c
+++ b/src/sort.c
@@ -2,6 +2,8 @@
 #include "pqsort.h" /* Partial qsort for SORT+LIMIT */
 #include <math.h> /* isnan() */
 
+zskiplistNode* zslGetElementByRank(zskiplist *zsl, unsigned long rank);
+
 redisSortOperation *createSortOperation(int type, robj *pattern) {
     redisSortOperation *so = zmalloc(sizeof(*so));
     so->type = type;
@@ -156,8 +158,9 @@ void sortCommand(redisClient *c) {
 
     /* Lookup the key to sort. It must be of the right types */
     sortval = lookupKeyRead(c->db,c->argv[1]);
-    if (sortval && sortval->type != REDIS_SET && sortval->type != REDIS_LIST &&
-        sortval->type != REDIS_ZSET)
+    if (sortval && sortval->type != REDIS_SET &&
+                   sortval->type != REDIS_LIST &&
+                   sortval->type != REDIS_ZSET)
     {
         addReply(c,shared.wrongtypeerr);
         return;
@@ -167,7 +170,7 @@ void sortCommand(redisClient *c) {
      * Operations can be GET/DEL/INCR/DECR */
     operations = listCreate();
     listSetFreeMethod(operations,zfree);
-    j = 2;
+    j = 2; /* options start at argv[2] */
 
     /* Now we need to protect sortval incrementing its count, in the future
      * SORT may have options able to overwrite/delete keys during the sorting
@@ -213,13 +216,18 @@ void sortCommand(redisClient *c) {
         j++;
     }
 
-    /* If we have STORE we need to force sorting for deterministic output
-     * and replication. We use alpha sorting since this is guaranteed to
-     * work with any input.
+    /* For the STORE option, or when SORT is called from a Lua script,
+     * we want to force a specific ordering even when no explicit ordering
+     * was asked (SORT BY nosort). This guarantees that replication / AOF
+     * is deterministic.
      *
-     * We also want determinism when SORT is called from Lua scripts, so
-     * in this case we also force alpha sorting. */
-    if ((storekey || c->flags & REDIS_LUA_CLIENT) && dontsort) {
+     * However in the case 'dontsort' is true, but the type to sort is a
+     * sorted set, we don't need to do anything as ordering is guaranteed
+     * in this special case. */
+    if ((storekey || c->flags & REDIS_LUA_CLIENT) &&
+        (dontsort && sortval->type != REDIS_ZSET))
+    {
+        /* Force ALPHA sorting */
         dontsort = 0;
         alpha = 1;
         sortby = NULL;
@@ -229,13 +237,41 @@ void sortCommand(redisClient *c) {
     if (sortval->type == REDIS_ZSET)
         zsetConvert(sortval, REDIS_ENCODING_SKIPLIST);
 
-    /* Load the sorting vector with all the objects to sort */
+    /* Objtain the length of the object to sort. */
     switch(sortval->type) {
     case REDIS_LIST: vectorlen = listTypeLength(sortval); break;
     case REDIS_SET: vectorlen =  setTypeSize(sortval); break;
     case REDIS_ZSET: vectorlen = dictSize(((zset*)sortval->ptr)->dict); break;
     default: vectorlen = 0; redisPanic("Bad SORT type"); /* Avoid GCC warning */
     }
+
+    /* Perform LIMIT start,count sanity checking. */
+    start = (limit_start < 0) ? 0 : limit_start;
+    end = (limit_count < 0) ? vectorlen-1 : start+limit_count-1;
+    if (start >= vectorlen) {
+        start = vectorlen-1;
+        end = vectorlen-2;
+    }
+    if (end >= vectorlen) end = vectorlen-1;
+
+    /* Optimization:
+     *
+     * 1) if the object to sort is a sorted set.
+     * 2) There is nothing to sort as dontsort is true (BY <constant string>).
+     * 3) We have a LIMIT option that actually reduces the number of elements
+     *    to fetch.
+     *
+     * In this case to load all the objects in the vector is a huge waste of
+     * resources. We just allocate a vector that is big enough for the selected
+     * range length, and make sure to load just this part in the vector. */
+    if (sortval->type == REDIS_ZSET &&
+        dontsort &&
+        (start != 0 || end != vectorlen-1))
+    {
+        vectorlen = end-start+1;
+    }
+
+    /* Load the sorting vector with all the objects to sort */
     vector = zmalloc(sizeof(redisSortObject)*vectorlen);
     j = 0;
 
@@ -259,6 +295,48 @@ void sortCommand(redisClient *c) {
             j++;
         }
         setTypeReleaseIterator(si);
+    } else if (sortval->type == REDIS_ZSET && dontsort) {
+        /* Special handling for a sorted set, if 'dontsort' is true.
+         * This makes sure we return elements in the sorted set original
+         * ordering, accordingly to DESC / ASC options.
+         *
+         * Note that in this case we also handle LIMIT here in a direct
+         * way, just getting the required range, as an optimization. */
+
+        zset *zs = sortval->ptr;
+        zskiplist *zsl = zs->zsl;
+        zskiplistNode *ln;
+        robj *ele;
+        int rangelen = vectorlen;
+
+        /* Check if starting point is trivial, before doing log(N) lookup. */
+        if (desc) {
+            long zsetlen = dictSize(((zset*)sortval->ptr)->dict);
+
+            ln = zsl->tail;
+            if (start > 0)
+                ln = zslGetElementByRank(zsl,zsetlen-start);
+        } else {
+            ln = zsl->header->level[0].forward;
+            if (start > 0)
+                ln = zslGetElementByRank(zsl,start+1);
+        }
+
+        while(rangelen--) {
+            redisAssertWithInfo(c,sortval,ln != NULL);
+            ele = ln->obj;
+            vector[j].obj = ele;
+            vector[j].u.score = 0;
+            vector[j].u.cmpobj = NULL;
+            j++;
+            ln = desc ? ln->backward : ln->level[0].forward;
+        }
+        /* The code producing the output does not know that in the case of
+         * sorted set, 'dontsort', and LIMIT, we are able to get just the
+         * range, already sorted, so we need to adjust "start" and "end"
+         * to make sure start is set to 0. */
+        end -= start;
+        start = 0;
     } else if (sortval->type == REDIS_ZSET) {
         dict *set = ((zset*)sortval->ptr)->dict;
         dictIterator *di;
@@ -319,16 +397,6 @@ void sortCommand(redisClient *c) {
         }
     }
 
-    /* We are ready to sort the vector... perform a bit of sanity check
-     * on the LIMIT option too. We'll use a partial version of quicksort. */
-    start = (limit_start < 0) ? 0 : limit_start;
-    end = (limit_count < 0) ? vectorlen-1 : start+limit_count-1;
-    if (start >= vectorlen) {
-        start = vectorlen-1;
-        end = vectorlen-2;
-    }
-    if (end >= vectorlen) end = vectorlen-1;
-
     if (dontsort == 0) {
         server.sort_desc = desc;
         server.sort_alpha = alpha;
diff --git a/tests/unit/sort.tcl b/tests/unit/sort.tcl
index 5a181641c..6c5644a79 100644
--- a/tests/unit/sort.tcl
+++ b/tests/unit/sort.tcl
@@ -118,6 +118,47 @@ start_server {
         r sort zset alpha desc
     } {e d c b a}
 
+    test "SORT sorted set BY nosort should retain ordering" {
+        r del zset
+        r zadd zset 1 a
+        r zadd zset 5 b
+        r zadd zset 2 c
+        r zadd zset 10 d
+        r zadd zset 3 e
+        r multi
+        r sort zset by nosort asc
+        r sort zset by nosort desc
+        r exec
+    } {{a c e b d} {d b e c a}}
+
+    test "SORT sorted set BY nosort + LIMIT" {
+        r del zset
+        r zadd zset 1 a
+        r zadd zset 5 b
+        r zadd zset 2 c
+        r zadd zset 10 d
+        r zadd zset 3 e
+        assert_equal [r sort zset by nosort asc limit 0 1] {a}
+        assert_equal [r sort zset by nosort desc limit 0 1] {d}
+        assert_equal [r sort zset by nosort asc limit 0 2] {a c}
+        assert_equal [r sort zset by nosort desc limit 0 2] {d b}
+        assert_equal [r sort zset by nosort limit 5 10] {}
+        assert_equal [r sort zset by nosort limit -10 100] {a c e b d}
+    }
+
+    test "SORT sorted set BY nosort works as expected from scripts" {
+        r del zset
+        r zadd zset 1 a
+        r zadd zset 5 b
+        r zadd zset 2 c
+        r zadd zset 10 d
+        r zadd zset 3 e
+        r eval {
+            return {redis.call('sort','zset','by','nosort','asc'),
+                    redis.call('sort','zset','by','nosort','desc')}
+        } 0
+    } {{a c e b d} {d b e c a}}
+
     test "SORT sorted set: +inf and -inf handling" {
         r del zset
         r zadd zset -100 a