From: antirez <antirez@gmail.com>
Date: Tue, 9 Mar 2010 15:25:55 +0000 (+0100)
Subject: Merged zsetops branch from Pietern
X-Git-Url: https://git.saurik.com/redis.git/commitdiff_plain/c1d57a106345a811c91a0bf1b6fff42ccbae61a3

Merged zsetops branch from Pietern
---

c1d57a106345a811c91a0bf1b6fff42ccbae61a3
diff --cc redis.c
index 8658bc19,cc64efb8..d15efbbc
--- a/redis.c
+++ b/redis.c
@@@ -674,10 -673,10 +674,12 @@@ static void brpopCommand(redisClient *c
  static void appendCommand(redisClient *c);
  static void substrCommand(redisClient *c);
  static void zrankCommand(redisClient *c);
 +static void zrevrankCommand(redisClient *c);
  static void hsetCommand(redisClient *c);
  static void hgetCommand(redisClient *c);
 +static void zremrangebyrankCommand(redisClient *c);
+ static void zunionCommand(redisClient *c);
+ static void zinterCommand(redisClient *c);
  
  /*================================= Globals ================================= */
  
@@@ -725,7 -724,8 +727,9 @@@ static struct redisCommand cmdTable[] 
      {"zincrby",zincrbyCommand,4,REDIS_CMD_BULK|REDIS_CMD_DENYOOM,1,1,1},
      {"zrem",zremCommand,3,REDIS_CMD_BULK,1,1,1},
      {"zremrangebyscore",zremrangebyscoreCommand,4,REDIS_CMD_INLINE,1,1,1},
 +    {"zremrangebyrank",zremrangebyrankCommand,4,REDIS_CMD_INLINE,1,1,1},
+     {"zunion",zunionCommand,-4,REDIS_CMD_INLINE|REDIS_CMD_DENYOOM,0,0,0},
+     {"zinter",zinterCommand,-4,REDIS_CMD_INLINE|REDIS_CMD_DENYOOM,0,0,0},
      {"zrange",zrangeCommand,-4,REDIS_CMD_INLINE,1,1,1},
      {"zrangebyscore",zrangebyscoreCommand,-4,REDIS_CMD_INLINE,1,1,1},
      {"zcount",zcountCommand,4,REDIS_CMD_INLINE,1,1,1},
@@@ -5410,47 -5330,171 +5415,212 @@@ static void zremrangebyscoreCommand(red
      }
  }
  
 +static void zremrangebyrankCommand(redisClient *c) {
 +    int start = atoi(c->argv[2]->ptr);
 +    int end = atoi(c->argv[3]->ptr);
 +    robj *zsetobj;
 +    zset *zs;
 +
 +    zsetobj = lookupKeyWrite(c->db,c->argv[1]);
 +    if (zsetobj == NULL) {
 +        addReply(c,shared.czero);
 +    } else {
 +        if (zsetobj->type != REDIS_ZSET) {
 +            addReply(c,shared.wrongtypeerr);
 +            return;
 +        }
 +
 +        zs = zsetobj->ptr;
 +        int llen = zs->zsl->length;
 +        long deleted;
 +
 +        /* convert negative indexes */
 +        if (start < 0) start = llen+start;
 +        if (end < 0) end = llen+end;
 +        if (start < 0) start = 0;
 +        if (end < 0) end = 0;
 +
 +        /* indexes sanity checks */
 +        if (start > end || start >= llen) {
 +            addReply(c,shared.czero);
 +            return;
 +        }
 +        if (end >= llen) end = llen-1;
 +
 +        /* increment start and end because zsl*Rank functions
 +         * use 1-based rank */
 +        deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
 +        if (htNeedsResize(zs->dict)) dictResize(zs->dict);
 +        server.dirty += deleted;
 +        addReplyLong(c, deleted);
 +    }
 +}
 +
+ typedef struct {
+     dict *dict;
+     double weight;
+ } zsetopsrc;
+ 
+ static int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
+     zsetopsrc *d1 = (void*) s1, *d2 = (void*) s2;
+     unsigned long size1, size2;
+     size1 = d1->dict ? dictSize(d1->dict) : 0;
+     size2 = d2->dict ? dictSize(d2->dict) : 0;
+     return size1 - size2;
+ }
+ 
+ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
+     int i, j, zsetnum;
+     zsetopsrc *src;
+     robj *dstobj;
+     zset *dstzset;
+     dictIterator *di;
+     dictEntry *de;
+ 
+     /* expect zsetnum input keys to be given */
+     zsetnum = atoi(c->argv[2]->ptr);
+     if (zsetnum < 1) {
+         addReplySds(c,sdsnew("-ERR at least 1 input key is needed for ZUNION/ZINTER\r\n"));
+         return;
+     }
+ 
+     /* test if the expected number of keys would overflow */
+     if (3+zsetnum > c->argc) {
+         addReply(c,shared.syntaxerr);
+         return;
+     }
+ 
+     /* read keys to be used for input */
+     src = malloc(sizeof(zsetopsrc) * zsetnum);
+     for (i = 0, j = 3; i < zsetnum; i++, j++) {
+         robj *zsetobj = lookupKeyWrite(c->db,c->argv[j]);
+         if (!zsetobj) {
+             src[i].dict = NULL;
+         } else {
+             if (zsetobj->type != REDIS_ZSET) {
+                 zfree(src);
+                 addReply(c,shared.wrongtypeerr);
+                 return;
+             }
+             src[i].dict = ((zset*)zsetobj->ptr)->dict;
+         }
+ 
+         /* default all weights to 1 */
+         src[i].weight = 1.0;
+     }
+ 
+     /* parse optional extra arguments */
+     if (j < c->argc) {
+         int remaining = c->argc-j;
+ 
+         while (remaining) {
+             if (!strcasecmp(c->argv[j]->ptr,"weights")) {
+                 j++; remaining--;
+                 if (remaining < zsetnum) {
+                     zfree(src);
+                     addReplySds(c,sdsnew("-ERR not enough weights for ZUNION/ZINTER\r\n"));
+                     return;
+                 }
+                 for (i = 0; i < zsetnum; i++, j++, remaining--) {
+                     src[i].weight = strtod(c->argv[j]->ptr, NULL);
+                 }
+             } else {
+                 zfree(src);
+                 addReply(c,shared.syntaxerr);
+                 return;
+             }
+         }
+     }
+ 
+     dstobj = createZsetObject();
+     dstzset = dstobj->ptr;
+ 
+     if (op == REDIS_OP_INTER) {
+         /* sort sets from the smallest to largest, this will improve our
+          * algorithm's performance */
+         qsort(src,zsetnum,sizeof(zsetopsrc), qsortCompareZsetopsrcByCardinality);
+ 
+         /* skip going over all entries if the smallest zset is NULL or empty */
+         if (src[0].dict && dictSize(src[0].dict) > 0) {
+             /* precondition: as src[0].dict is non-empty and the zsets are ordered
+              * from small to large, all src[i > 0].dict are non-empty too */
+             di = dictGetIterator(src[0].dict);
+             while((de = dictNext(di)) != NULL) {
+                 double *score = zmalloc(sizeof(double));
+                 *score = 0.0;
+ 
+                 for (j = 0; j < zsetnum; j++) {
+                     dictEntry *other = (j == 0) ? de : dictFind(src[j].dict,dictGetEntryKey(de));
+                     if (other) {
+                         *score = *score + src[j].weight * (*(double*)dictGetEntryVal(other));
+                     } else {
+                         break;
+                     }
+                 }
+ 
+                 /* skip entry when not present in every source dict */
+                 if (j != zsetnum) {
+                     zfree(score);
+                 } else {
+                     robj *o = dictGetEntryKey(de);
+                     dictAdd(dstzset->dict,o,score);
+                     incrRefCount(o); /* added to dictionary */
+                     zslInsert(dstzset->zsl,*score,o);
+                     incrRefCount(o); /* added to skiplist */
+                 }
+             }
+             dictReleaseIterator(di);
+         }
+     } else if (op == REDIS_OP_UNION) {
+         for (i = 0; i < zsetnum; i++) {
+             if (!src[i].dict) continue;
+ 
+             di = dictGetIterator(src[i].dict);
+             while((de = dictNext(di)) != NULL) {
+                 /* skip key when already processed */
+                 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) continue;
+ 
+                 double *score = zmalloc(sizeof(double));
+                 *score = 0.0;
+                 for (j = 0; j < zsetnum; j++) {
+                     if (!src[j].dict) continue;
+ 
+                     dictEntry *other = (i == j) ? de : dictFind(src[j].dict,dictGetEntryKey(de));
+                     if (other) {
+                         *score = *score + src[j].weight * (*(double*)dictGetEntryVal(other));
+                     }
+                 }
+ 
+                 robj *o = dictGetEntryKey(de);
+                 dictAdd(dstzset->dict,o,score);
+                 incrRefCount(o); /* added to dictionary */
+                 zslInsert(dstzset->zsl,*score,o);
+                 incrRefCount(o); /* added to skiplist */
+             }
+             dictReleaseIterator(di);
+         }
+     } else {
+         /* unknown operator */
+         redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION);
+     }
+ 
+     deleteKey(c->db,dstkey);
+     dictAdd(c->db->dict,dstkey,dstobj);
+     incrRefCount(dstkey);
+ 
+     addReplyLong(c, dstzset->zsl->length);
+     server.dirty++;
+     zfree(src);
+ }
+ 
+ static void zunionCommand(redisClient *c) {
+     zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION);
+ }
+ 
+ static void zinterCommand(redisClient *c) {
+     zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
+ }
+ 
  static void zrangeGenericCommand(redisClient *c, int reverse) {
      robj *o;
      int start = atoi(c->argv[2]->ptr);
diff --cc test-redis.tcl
index 9139f544,54b614fd..00370a4c
--- a/test-redis.tcl
+++ b/test-redis.tcl
@@@ -1466,16 -1462,29 +1466,39 @@@ proc main {server port} 
          list [$r zremrangebyscore zset -inf +inf] [$r zrange zset 0 -1]
      } {5 {}}
  
 +    test {ZREMRANGEBYRANK basics} {
 +        $r del zset
 +        $r zadd zset 1 a
 +        $r zadd zset 2 b
 +        $r zadd zset 3 c
 +        $r zadd zset 4 d
 +        $r zadd zset 5 e
 +        list [$r zremrangebyrank zset 1 3] [$r zrange zset 0 -1]
 +    } {3 {a e}}
 +
+     test {ZUNION basics} {
+         $r del zseta zsetb zsetc
+         $r zadd zseta 1 a
+         $r zadd zseta 2 b
+         $r zadd zseta 3 c
+         $r zadd zsetb 1 b
+         $r zadd zsetb 2 c
+         $r zadd zsetb 3 d
+         list [$r zunion zsetc 2 zseta zsetb] [$r zrange zsetc 0 -1 withscores]
+     } {4 {a 1 b 3 d 3 c 5}}
+ 
+     test {ZUNION with weights} {
+         list [$r zunion zsetc 2 zseta zsetb weights 2 3] [$r zrange zsetc 0 -1 withscores]
+     } {4 {a 2 b 7 d 9 c 12}}
+ 
+     test {ZINTER basics} {
+         list [$r zinter zsetc 2 zseta zsetb] [$r zrange zsetc 0 -1 withscores]
+     } {2 {b 3 c 5}}
+ 
+     test {ZINTER with weights} {
+         list [$r zinter zsetc 2 zseta zsetb weights 2 3] [$r zrange zsetc 0 -1 withscores]
+     } {2 {b 7 c 12}}
+ 
      test {SORT against sorted sets} {
          $r del zset
          $r zadd zset 1 a