Skip to content
This repository has been archived by the owner on Jul 11, 2021. It is now read-only.

Commit

Permalink
Add ZRANK UNIQUE
Browse files Browse the repository at this point in the history 
Includes tests

Closes redis#943
  • Loading branch information
@mattsta
Sylvain Royer authored and mattsta committed Dec 11, 2014
1 parent 3cd36a4 commit 9605949
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 27 deletions.
4 changes: 2 additions & 2 deletions src/redis.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -189,8 +189,8 @@ struct redisCommand redisCommandTable[] = {
{"zrevrange",zrevrangeCommand,-4,"r",0,NULL,1,1,1,0,0},
{"zcard",zcardCommand,2,"rF",0,NULL,1,1,1,0,0},
{"zscore",zscoreCommand,3,"rF",0,NULL,1,1,1,0,0},
{"zrank",zrankCommand,3,"rF",0,NULL,1,1,1,0,0},
{"zrevrank",zrevrankCommand,3,"rF",0,NULL,1,1,1,0,0},
{"zrank",zrankCommand,-3,"rF",0,NULL,1,1,1,0,0},
{"zrevrank",zrevrankCommand,-3,"rF",0,NULL,1,1,1,0,0},
{"zscan",zscanCommand,-3,"rR",0,NULL,1,1,1,0,0},
{"hset",hsetCommand,4,"wmF",0,NULL,1,1,1,0,0},
{"hsetnx",hsetnxCommand,4,"wmF",0,NULL,1,1,1,0,0},
Expand Down
80 changes: 58 additions & 22 deletions src/t_zset.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -383,23 +383,33 @@ unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned
* Returns 0 when the element cannot be found, rank otherwise.
* Note that the rank is 1-based due to the span of zsl->header to the
* first element. */
unsigned long zslGetRank(zskiplist *zsl, double score, robj *o) {
zskiplistNode *x;
unsigned long zslGetRank(zskiplist *zsl, double score, robj *o, int unique, int reverse) {
zskiplistNode *x, *nextNode;
unsigned long rank = 0;
int i;

x = zsl->header;
for (i = zsl->level-1; i >= 0; i--) {
while (x->level[i].forward &&
(x->level[i].forward->score < score ||
(x->level[i].forward->score == score &&
compareStringObjects(x->level[i].forward->obj,o) <= 0))) {
nextNode = x->level[i].forward;
while (nextNode &&
(nextNode->score < score ||
(nextNode->score == score &&
((!unique && compareStringObjects(nextNode->obj,o) <= 0) || // Check for ZRANK
(unique && (!nextNode->backward || nextNode->backward->score < score)) || // Check for ZRANK UNIQUE
(unique && reverse) // Check for ZREVRANK UNIQUE
))))
{
rank += x->level[i].span;
x = x->level[i].forward;
x = nextNode;
nextNode=x->level[i].forward;
}

/* x might be equal to zsl->header, so test if obj is non-NULL */
if (x->obj && equalStringObjects(x->obj,o)) {
if (x->obj && x->score == score &&
((!unique && equalStringObjects(x->obj,o)) || // ZRANK check.
(!reverse && unique && (!x->backward || x->backward->score < score)) || // ZRANK UNIQUE check
(reverse && unique && (!x->level[0].forward || x->level[0].forward->score > score)))) // ZREVRANK UNIQUE check
{
return rank;
}
}
Expand Down Expand Up @@ -2530,15 +2540,15 @@ void zlexcountCommand(redisClient *c) {

/* Use rank of first element, if any, to determine preliminary count */
if (zn != NULL) {
rank = zslGetRank(zsl, zn->score, zn->obj);
rank = zslGetRank(zsl, zn->score, zn->obj, 0, 0);
count = (zsl->length - (rank - 1));

/* Find last element in range */
zn = zslLastInLexRange(zsl, &range);

/* Use rank of last element, if any, to determine the actual count */
if (zn != NULL) {
rank = zslGetRank(zsl, zn->score, zn->obj);
rank = zslGetRank(zsl, zn->score, zn->obj, 0, 0);
count -= (zsl->length - rank);
}
}
Expand Down Expand Up @@ -2779,9 +2789,17 @@ void zrankGenericCommand(redisClient *c, int reverse) {
robj *key = c->argv[1];
robj *ele = c->argv[2];
robj *zobj;
int unique = 0;
unsigned long llen;
unsigned long rank;

if (c->argc == 4 && !strcasecmp(c->argv[3]->ptr,"unique")) {
unique = 1;
} else if (c->argc >= 4) {
addReply(c,shared.syntaxerr);
return;
}

if ((zobj = lookupKeyReadOrReply(c,key,shared.nullbulk)) == NULL ||
checkType(c,zobj,REDIS_ZSET)) return;
llen = zsetLength(zobj);
Expand All @@ -2791,25 +2809,43 @@ void zrankGenericCommand(redisClient *c, int reverse) {
if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
unsigned char *zl = zobj->ptr;
unsigned char *eptr, *sptr;
double currentScore, prevScore, count;

eptr = ziplistIndex(zl,0);
redisAssertWithInfo(c,zobj,eptr != NULL);
sptr = ziplistNext(zl,eptr);
redisAssertWithInfo(c,zobj,sptr != NULL);

rank = 1;
if (reverse) {
sptr = ziplistIndex(zl,-1);
redisAssertWithInfo(c,zobj,sptr != NULL);
eptr = ziplistPrev(zl,sptr);
redisAssertWithInfo(c,zobj,eptr != NULL);
}
else {
eptr = ziplistIndex(zl,0);
redisAssertWithInfo(c,zobj,eptr != NULL);
sptr = ziplistNext(zl,eptr);
redisAssertWithInfo(c,zobj,sptr != NULL);
}

count = 1;
rank = count;
prevScore = 0;

while(eptr != NULL) {
currentScore = zzlGetScore(sptr);

if ((!unique) || (prevScore != currentScore)) {
prevScore = currentScore;
rank = count;
}

if (ziplistCompare(eptr,ele->ptr,sdslen(ele->ptr)))
break;
rank++;
zzlNext(zl,&eptr,&sptr);

count++;
reverse ? zzlPrev(zl,&eptr,&sptr) : zzlNext(zl,&eptr,&sptr);
}

if (eptr != NULL) {
if (reverse)
addReplyLongLong(c,llen-rank);
else
addReplyLongLong(c,rank-1);
addReplyLongLong(c,rank-1);
} else {
addReply(c,shared.nullbulk);
}
Expand All @@ -2823,7 +2859,7 @@ void zrankGenericCommand(redisClient *c, int reverse) {
de = dictFind(zs->dict,ele);
if (de != NULL) {
score = *(double*)dictGetVal(de);
rank = zslGetRank(zsl,score,ele);
rank = zslGetRank(zsl,score,ele,unique, reverse);
redisAssertWithInfo(c,ele,rank); /* Existing elements always have a rank. */
if (reverse)
addReplyLongLong(c,llen-rank);
Expand Down
53 changes: 50 additions & 3 deletions tests/unit/type/zset.tcl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -168,21 +168,68 @@ start_server {tags {"zset"}} {
assert_equal {d 4 c 3 b 2 a 1} [r zrevrange ztmp 0 -1 withscores]
}

test "ZRANK/ZREVRANK basics - $encoding" {
test "ZRANK basics - $encoding" {
r del zranktmp
r zadd zranktmp 10 x
r zadd zranktmp 20 y
r zadd zranktmp -10 x
r zadd zranktmp 0 y
r zadd zranktmp 30 z
assert_equal 0 [r zrank zranktmp x]
assert_equal 1 [r zrank zranktmp y]
assert_equal 2 [r zrank zranktmp z]
assert_equal "" [r zrank zranktmp foo]
}

test "ZREVRANK basics - $encoding" {
r del zranktmp
r zadd zranktmp -10 x
r zadd zranktmp 0 y
r zadd zranktmp 30 z
assert_equal 2 [r zrevrank zranktmp x]
assert_equal 1 [r zrevrank zranktmp y]
assert_equal 0 [r zrevrank zranktmp z]
assert_equal "" [r zrevrank zranktmp foo]
}

test "ZRANK negative score - $encoding" {
r del zranktmp
r zadd zranktmp -1 x
assert_equal 0 [r zrank zranktmp x]
}

test "ZRANK unique - $encoding" {
r del zranktmp
r zadd zranktmp -10 x
r zadd zranktmp 0 y
r zadd zranktmp 0 z
r zadd zranktmp 30 d
r zadd zranktmp 30 e
r zadd zranktmp 40 f
assert_equal 0 [r zrank zranktmp x unique]
assert_equal 1 [r zrank zranktmp y unique]
assert_equal 1 [r zrank zranktmp z unique]
assert_equal 3 [r zrank zranktmp d unique]
assert_equal 3 [r zrank zranktmp e unique]
assert_equal 5 [r zrank zranktmp f unique]
assert_equal "" [r zrank zranktmp foo unique]
}

test "ZREVRANK unique - $encoding" {
r del zranktmp
r zadd zranktmp -10 x
r zadd zranktmp 0 y
r zadd zranktmp 0 z
r zadd zranktmp 30 d
r zadd zranktmp 30 e
r zadd zranktmp 40 f
assert_equal 5 [r zrevrank zranktmp x unique]
assert_equal 3 [r zrevrank zranktmp y unique]
assert_equal 3 [r zrevrank zranktmp z unique]
assert_equal 1 [r zrevrank zranktmp d unique]
assert_equal 1 [r zrevrank zranktmp e unique]
assert_equal 0 [r zrevrank zranktmp f unique]
assert_equal "" [r zrank zranktmp foo unique]
}

test "ZRANK - after deletion - $encoding" {
r zrem zranktmp y
assert_equal 0 [r zrank zranktmp x]
Expand Down

0 comments on commit 9605949

Please sign in to comment.