Variadic ZADD

This commit is contained in:
antirez 2011-05-31 17:47:34 +02:00
parent 632e4c09ac
commit ef231a7c56
2 changed files with 97 additions and 85 deletions

View File

@ -116,7 +116,7 @@ struct redisCommand redisCommandTable[] = {
{"sdiff",sdiffCommand,-2,REDIS_CMD_DENYOOM,NULL,1,-1,1,0,0}, {"sdiff",sdiffCommand,-2,REDIS_CMD_DENYOOM,NULL,1,-1,1,0,0},
{"sdiffstore",sdiffstoreCommand,-3,REDIS_CMD_DENYOOM,NULL,2,-1,1,0,0}, {"sdiffstore",sdiffstoreCommand,-3,REDIS_CMD_DENYOOM,NULL,2,-1,1,0,0},
{"smembers",sinterCommand,2,0,NULL,1,1,1,0,0}, {"smembers",sinterCommand,2,0,NULL,1,1,1,0,0},
{"zadd",zaddCommand,4,REDIS_CMD_DENYOOM,NULL,1,1,1,0,0}, {"zadd",zaddCommand,-4,REDIS_CMD_DENYOOM,NULL,1,1,1,0,0},
{"zincrby",zincrbyCommand,4,REDIS_CMD_DENYOOM,NULL,1,1,1,0,0}, {"zincrby",zincrbyCommand,4,REDIS_CMD_DENYOOM,NULL,1,1,1,0,0},
{"zrem",zremCommand,3,0,NULL,1,1,1,0,0}, {"zrem",zremCommand,3,0,NULL,1,1,1,0,0},
{"zremrangebyscore",zremrangebyscoreCommand,4,0,NULL,1,1,1,0,0}, {"zremrangebyscore",zremrangebyscoreCommand,4,0,NULL,1,1,1,0,0},

View File

@ -810,11 +810,29 @@ void zaddGenericCommand(redisClient *c, int incr) {
robj *ele; robj *ele;
robj *zobj; robj *zobj;
robj *curobj; robj *curobj;
double score, curscore = 0.0; double score = 0, *scores, curscore = 0.0;
int j, elements = (c->argc-2)/2;
int added = 0;
if (getDoubleFromObjectOrReply(c,c->argv[2],&score,NULL) != REDIS_OK) if (c->argc % 2) {
addReply(c,shared.syntaxerr);
return; return;
}
/* Start parsing all the scores, we need to emit any syntax error
* before executing additions to the sorted set, as the command should
* either execute fully or nothing at all. */
scores = zmalloc(sizeof(double)*elements);
for (j = 0; j < elements; j++) {
if (getDoubleFromObjectOrReply(c,c->argv[2+j*2],&scores[j],NULL)
!= REDIS_OK)
{
zfree(scores);
return;
}
}
/* Lookup the key and create the sorted set if does not exist. */
zobj = lookupKeyWrite(c->db,key); zobj = lookupKeyWrite(c->db,key);
if (zobj == NULL) { if (zobj == NULL) {
if (server.zset_max_ziplist_entries == 0 || if (server.zset_max_ziplist_entries == 0 ||
@ -828,111 +846,105 @@ void zaddGenericCommand(redisClient *c, int incr) {
} else { } else {
if (zobj->type != REDIS_ZSET) { if (zobj->type != REDIS_ZSET) {
addReply(c,shared.wrongtypeerr); addReply(c,shared.wrongtypeerr);
zfree(scores);
return; return;
} }
} }
if (zobj->encoding == REDIS_ENCODING_ZIPLIST) { for (j = 0; j < elements; j++) {
unsigned char *eptr; score = scores[j];
/* Prefer non-encoded element when dealing with ziplists. */ if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
ele = c->argv[3]; unsigned char *eptr;
if ((eptr = zzlFind(zobj->ptr,ele,&curscore)) != NULL) {
if (incr) { /* Prefer non-encoded element when dealing with ziplists. */
score += curscore; ele = c->argv[3+j*2];
if (isnan(score)) { if ((eptr = zzlFind(zobj->ptr,ele,&curscore)) != NULL) {
addReplyError(c,nanerr); if (incr) {
/* Don't need to check if the sorted set is empty, because score += curscore;
* we know it has at least one element. */ if (isnan(score)) {
return; addReplyError(c,nanerr);
/* Don't need to check if the sorted set is empty
* because we know it has at least one element. */
zfree(scores);
return;
}
} }
}
/* Remove and re-insert when score changed. */ /* Remove and re-insert when score changed. */
if (score != curscore) { if (score != curscore) {
zobj->ptr = zzlDelete(zobj->ptr,eptr); zobj->ptr = zzlDelete(zobj->ptr,eptr);
zobj->ptr = zzlInsert(zobj->ptr,ele,score);
signalModifiedKey(c->db,key);
server.dirty++;
}
} else {
/* Optimize: check if the element is too large or the list
* becomes too long *before* executing zzlInsert. */
zobj->ptr = zzlInsert(zobj->ptr,ele,score); zobj->ptr = zzlInsert(zobj->ptr,ele,score);
if (zzlLength(zobj->ptr) > server.zset_max_ziplist_entries)
zsetConvert(zobj,REDIS_ENCODING_SKIPLIST);
if (sdslen(ele->ptr) > server.zset_max_ziplist_value)
zsetConvert(zobj,REDIS_ENCODING_SKIPLIST);
signalModifiedKey(c->db,key); signalModifiedKey(c->db,key);
server.dirty++; server.dirty++;
if (!incr) added++;
} }
} else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) {
zset *zs = zobj->ptr;
zskiplistNode *znode;
dictEntry *de;
if (incr) /* ZINCRBY */ ele = c->argv[3+j*2] = tryObjectEncoding(c->argv[3+j*2]);
addReplyDouble(c,score); de = dictFind(zs->dict,ele);
else /* ZADD */ if (de != NULL) {
addReply(c,shared.czero); curobj = dictGetEntryKey(de);
} else { curscore = *(double*)dictGetEntryVal(de);
/* Optimize: check if the element is too large or the list becomes
* too long *before* executing zzlInsert. */
zobj->ptr = zzlInsert(zobj->ptr,ele,score);
if (zzlLength(zobj->ptr) > server.zset_max_ziplist_entries)
zsetConvert(zobj,REDIS_ENCODING_SKIPLIST);
if (sdslen(ele->ptr) > server.zset_max_ziplist_value)
zsetConvert(zobj,REDIS_ENCODING_SKIPLIST);
signalModifiedKey(c->db,key); if (incr) {
server.dirty++; score += curscore;
if (isnan(score)) {
if (incr) /* ZINCRBY */ addReplyError(c,nanerr);
addReplyDouble(c,score); /* Don't need to check if the sorted set is empty
else /* ZADD */ * because we know it has at least one element. */
addReply(c,shared.cone); zfree(scores);
} return;
} else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) { }
zset *zs = zobj->ptr;
zskiplistNode *znode;
dictEntry *de;
ele = c->argv[3] = tryObjectEncoding(c->argv[3]);
de = dictFind(zs->dict,ele);
if (de != NULL) {
curobj = dictGetEntryKey(de);
curscore = *(double*)dictGetEntryVal(de);
if (incr) {
score += curscore;
if (isnan(score)) {
addReplyError(c,nanerr);
/* Don't need to check if the sorted set is empty, because
* we know it has at least one element. */
return;
} }
}
/* Remove and re-insert when score changed. We can safely delete /* Remove and re-insert when score changed. We can safely
* the key object from the skiplist, since the dictionary still has * delete the key object from the skiplist, since the
* a reference to it. */ * dictionary still has a reference to it. */
if (score != curscore) { if (score != curscore) {
redisAssert(zslDelete(zs->zsl,curscore,curobj)); redisAssert(zslDelete(zs->zsl,curscore,curobj));
znode = zslInsert(zs->zsl,score,curobj); znode = zslInsert(zs->zsl,score,curobj);
incrRefCount(curobj); /* Re-inserted in skiplist. */ incrRefCount(curobj); /* Re-inserted in skiplist. */
dictGetEntryVal(de) = &znode->score; /* Update score ptr. */ dictGetEntryVal(de) = &znode->score; /* Update score ptr. */
signalModifiedKey(c->db,key);
server.dirty++;
}
} else {
znode = zslInsert(zs->zsl,score,ele);
incrRefCount(ele); /* Inserted in skiplist. */
redisAssert(dictAdd(zs->dict,ele,&znode->score) == DICT_OK);
incrRefCount(ele); /* Added to dictionary. */
signalModifiedKey(c->db,key); signalModifiedKey(c->db,key);
server.dirty++; server.dirty++;
if (!incr) added++;
} }
if (incr) /* ZINCRBY */
addReplyDouble(c,score);
else /* ZADD */
addReply(c,shared.czero);
} else { } else {
znode = zslInsert(zs->zsl,score,ele); redisPanic("Unknown sorted set encoding");
incrRefCount(ele); /* Inserted in skiplist. */
redisAssert(dictAdd(zs->dict,ele,&znode->score) == DICT_OK);
incrRefCount(ele); /* Added to dictionary. */
signalModifiedKey(c->db,key);
server.dirty++;
if (incr) /* ZINCRBY */
addReplyDouble(c,score);
else /* ZADD */
addReply(c,shared.cone);
} }
} else {
redisPanic("Unknown sorted set encoding");
} }
zfree(scores);
if (incr) /* ZINCRBY */
addReplyDouble(c,score);
else /* ZADD */
addReplyLongLong(c,added);
} }
void zaddCommand(redisClient *c) { void zaddCommand(redisClient *c) {