use a struct to store both a dict and its weight for ZUNION and ZINTER, so qsort can be applied

This commit is contained in:
Pieter Noordhuis 2010-03-09 16:12:34 +01:00
parent 2830ca539b
commit 8f92e76831

81
redis.c
View File

@ -5330,10 +5330,22 @@ static void zremrangebyscoreCommand(redisClient *c) {
} }
} }
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) { static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
int i, j, k, zsetnum; int i, j, zsetnum;
dict **srcdicts; zsetopsrc *src;
double *weights;
robj *dstobj; robj *dstobj;
zset *dstzset; zset *dstzset;
dictIterator *di; dictIterator *di;
@ -5353,24 +5365,22 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
} }
/* read keys to be used for input */ /* read keys to be used for input */
srcdicts = zmalloc(sizeof(dict*) * zsetnum); src = malloc(sizeof(zsetopsrc) * zsetnum);
weights = zmalloc(sizeof(double) * zsetnum);
for (i = 0, j = 3; i < zsetnum; i++, j++) { for (i = 0, j = 3; i < zsetnum; i++, j++) {
robj *zsetobj = lookupKeyWrite(c->db,c->argv[j]); robj *zsetobj = lookupKeyWrite(c->db,c->argv[j]);
if (!zsetobj) { if (!zsetobj) {
srcdicts[i] = NULL; src[i].dict = NULL;
} else { } else {
if (zsetobj->type != REDIS_ZSET) { if (zsetobj->type != REDIS_ZSET) {
zfree(srcdicts); zfree(src);
zfree(weights);
addReply(c,shared.wrongtypeerr); addReply(c,shared.wrongtypeerr);
return; return;
} }
srcdicts[i] = ((zset*)zsetobj->ptr)->dict; src[i].dict = ((zset*)zsetobj->ptr)->dict;
} }
/* default all weights to 1 */ /* default all weights to 1 */
weights[i] = 1.0; src[i].weight = 1.0;
} }
/* parse optional extra arguments */ /* parse optional extra arguments */
@ -5381,17 +5391,15 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
if (!strcasecmp(c->argv[j]->ptr,"weights")) { if (!strcasecmp(c->argv[j]->ptr,"weights")) {
j++; remaining--; j++; remaining--;
if (remaining < zsetnum) { if (remaining < zsetnum) {
zfree(srcdicts); zfree(src);
zfree(weights);
addReplySds(c,sdsnew("-ERR not enough weights for ZUNION/ZINTER\r\n")); addReplySds(c,sdsnew("-ERR not enough weights for ZUNION/ZINTER\r\n"));
return; return;
} }
for (i = 0; i < zsetnum; i++, j++, remaining--) { for (i = 0; i < zsetnum; i++, j++, remaining--) {
weights[i] = strtod(c->argv[j]->ptr, NULL); src[i].weight = strtod(c->argv[j]->ptr, NULL);
} }
} else { } else {
zfree(srcdicts); zfree(src);
zfree(weights);
addReply(c,shared.syntaxerr); addReply(c,shared.syntaxerr);
return; return;
} }
@ -5402,34 +5410,30 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
dstzset = dstobj->ptr; dstzset = dstobj->ptr;
if (op == REDIS_OP_INTER) { if (op == REDIS_OP_INTER) {
/* store index of smallest zset in variable j */ /* sort sets from the smallest to largest, this will improve our
for (i = 0, j = 0; i < zsetnum; i++) { * algorithm's performance */
if (!srcdicts[i] || dictSize(srcdicts[i]) == 0) { qsort(src,zsetnum,sizeof(zsetopsrc), qsortCompareZsetopsrcByCardinality);
break;
} /* skip going over all entries if the smallest zset is NULL or empty */
if (dictSize(srcdicts[i]) < dictSize(srcdicts[j])) { if (src[0].dict && dictSize(src[0].dict) > 0) {
j = i; /* 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);
/* skip going over all entries if at least one dict was NULL or empty */
if (i == zsetnum) {
/* precondition: all srcdicts are non-NULL and non-empty */
di = dictGetIterator(srcdicts[j]);
while((de = dictNext(di)) != NULL) { while((de = dictNext(di)) != NULL) {
double *score = zmalloc(sizeof(double)); double *score = zmalloc(sizeof(double));
*score = 0.0; *score = 0.0;
for (k = 0; k < zsetnum; k++) { for (j = 0; j < zsetnum; j++) {
dictEntry *other = (k == j) ? de : dictFind(srcdicts[k],dictGetEntryKey(de)); dictEntry *other = (j == 0) ? de : dictFind(src[j].dict,dictGetEntryKey(de));
if (other) { if (other) {
*score = *score + weights[k] * (*(double*)dictGetEntryVal(other)); *score = *score + src[j].weight * (*(double*)dictGetEntryVal(other));
} else { } else {
break; break;
} }
} }
/* skip entry when not present in every source dict */ /* skip entry when not present in every source dict */
if (k != zsetnum) { if (j != zsetnum) {
zfree(score); zfree(score);
} else { } else {
robj *o = dictGetEntryKey(de); robj *o = dictGetEntryKey(de);
@ -5443,9 +5447,9 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
} }
} else if (op == REDIS_OP_UNION) { } else if (op == REDIS_OP_UNION) {
for (i = 0; i < zsetnum; i++) { for (i = 0; i < zsetnum; i++) {
if (!srcdicts[i]) continue; if (!src[i].dict) continue;
di = dictGetIterator(srcdicts[i]); di = dictGetIterator(src[i].dict);
while((de = dictNext(di)) != NULL) { while((de = dictNext(di)) != NULL) {
/* skip key when already processed */ /* skip key when already processed */
if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) continue; if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) continue;
@ -5453,11 +5457,11 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
double *score = zmalloc(sizeof(double)); double *score = zmalloc(sizeof(double));
*score = 0.0; *score = 0.0;
for (j = 0; j < zsetnum; j++) { for (j = 0; j < zsetnum; j++) {
if (!srcdicts[j]) continue; if (!src[j].dict) continue;
dictEntry *other = (i == j) ? de : dictFind(srcdicts[j],dictGetEntryKey(de)); dictEntry *other = (i == j) ? de : dictFind(src[j].dict,dictGetEntryKey(de));
if (other) { if (other) {
*score = *score + weights[j] * (*(double*)dictGetEntryVal(other)); *score = *score + src[j].weight * (*(double*)dictGetEntryVal(other));
} }
} }
@ -5480,8 +5484,7 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
addReplyLong(c, dstzset->zsl->length); addReplyLong(c, dstzset->zsl->length);
server.dirty++; server.dirty++;
zfree(srcdicts); zfree(src);
zfree(weights);
} }
static void zunionCommand(redisClient *c) { static void zunionCommand(redisClient *c) {