Permalink
Browse files

Added ZSET command ZDIFFSTORE

  • Loading branch information...
Michael Parrish
Michael Parrish committed Apr 10, 2012
1 parent 0b913c6 commit 6e78744506b5a0382ff458f590b97cc914e7051d
Showing with 58 additions and 3 deletions.
  1. +1 −0 src/redis.c
  2. +1 −0 src/redis.h
  3. +41 −3 src/t_zset.c
  4. +15 −0 tests/unit/type/zset.tcl
View
@@ -166,6 +166,7 @@ struct redisCommand redisCommandTable[] = {
{"zremrangebyrank",zremrangebyrankCommand,4,"w",0,NULL,1,1,1,0,0},
{"zunionstore",zunionstoreCommand,-4,"wm",0,zunionInterGetKeys,0,0,0,0,0},
{"zinterstore",zinterstoreCommand,-4,"wm",0,zunionInterGetKeys,0,0,0,0,0},
+ {"zdiffstore",zdiffstoreCommand,-4,"wm",0,zunionInterGetKeys,0,0,0,0,0},
{"zrange",zrangeCommand,-4,"r",0,NULL,1,1,1,0,0},
{"zrangebyscore",zrangebyscoreCommand,-4,"r",0,NULL,1,1,1,0,0},
{"zrevrangebyscore",zrevrangebyscoreCommand,-4,"r",0,NULL,1,1,1,0,0},
View
@@ -1219,6 +1219,7 @@ void hlenCommand(redisClient *c);
void zremrangebyrankCommand(redisClient *c);
void zunionstoreCommand(redisClient *c);
void zinterstoreCommand(redisClient *c);
+void zdiffstoreCommand(redisClient *c);
void hkeysCommand(redisClient *c);
void hvalsCommand(redisClient *c);
void hgetallCommand(redisClient *c);
View
@@ -1460,7 +1460,7 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
if (setnum < 1) {
addReplyError(c,
- "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE");
+ "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE/ZDIFFSTORE");
return;
}
@@ -1507,7 +1507,7 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
return;
}
}
- } else if (remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
+ } else if ( op != REDIS_OP_DIFF && remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
j++; remaining--;
if (!strcasecmp(c->argv[j]->ptr,"sum")) {
aggregate = REDIS_AGGR_SUM;
@@ -1534,7 +1534,9 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
/* sort sets from the smallest to largest, this will improve our
* algorithm's performance */
- qsort(src,setnum,sizeof(zsetopsrc),zuiCompareByCardinality);
+ if ( op != REDIS_OP_DIFF ) {
+ qsort(src, setnum, sizeof(zsetopsrc), zuiCompareByCardinality);
+ }
dstobj = createZsetObject();
dstzset = dstobj->ptr;
@@ -1620,6 +1622,38 @@ void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
maxelelen = sdslen(tmp->ptr);
}
}
+ } else if (op == REDIS_OP_DIFF) {
+ /* Skip everything if the smallest input is empty. */
+ if (zuiLength(&src[0]) > 0) {
+ while (zuiNext(&src[0],&zval)) {
+ double score, value;
+ int exists = 0;
+
+ // weighting doesn't make sense here
+ score = zval.score;
+ if (isnan(score)) score = 0;
+
+ // if this key exists in any other input set, stop looking for it
+ for (j = 1; j < setnum; j++) {
+ if (zuiFind(&src[j],&zval,&value)) {
+ exists = 1;
+ break;
+ }
+ }
+
+ if (!exists) {
+ tmp = zuiObjectFromValue(&zval);
+ znode = zslInsert(dstzset->zsl,score,tmp);
+ incrRefCount(zval.ele); /* added to skiplist */
+ dictAdd(dstzset->dict,tmp,&znode->score);
+ incrRefCount(zval.ele); /* added to dictionary */
+
+ if (tmp->encoding == REDIS_ENCODING_RAW)
+ if (sdslen(tmp->ptr) > maxelelen)
+ maxelelen = sdslen(tmp->ptr);
+ }
+ }
+ }
} else {
redisPanic("Unknown operator");
}
@@ -1657,6 +1691,10 @@ void zinterstoreCommand(redisClient *c) {
zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
}
+void zdiffstoreCommand(redisClient *c) {
+ zunionInterGenericCommand(c, c->argv[1], REDIS_OP_DIFF);
+}
+
void zrangeGenericCommand(redisClient *c, int reverse) {
robj *key = c->argv[1];
robj *zobj;
View
@@ -471,6 +471,21 @@ start_server {tags {"zset"}} {
assert_equal {b 2 c 3} [r zrange zsetc 0 -1 withscores]
}
+ test "ZDIFFSTORE basics - $encoding" {
+ assert_equal 1 [r zdiffstore zsetc 2 zseta zsetb]
+ assert_equal {a 1} [r zrange zsetc 0 -1 withscores]
+ }
+
+ test "ZDIFFSTORE with WITHSCORES - $encoding" {
+ r del zsetc
+ r zadd zseta 4 e
+ r zadd zseta 5 f
+ r zadd zsetc 4 e
+
+ assert_equal 2 [r zdiffstore zsetd 3 zseta zsetb zsetc]
+ assert_equal {a 1 f 5} [r zrange zsetd 0 -1 withscores]
+ }
+
foreach cmd {ZUNIONSTORE ZINTERSTORE} {
test "$cmd with +inf/-inf scores - $encoding" {
r del zsetinf1 zsetinf2

0 comments on commit 6e78744

Please sign in to comment.