allow regular sets to be passed to zunionstore/zinterstore

This commit is contained in:
Josiah Carlson 2010-05-25 17:41:35 -07:00 committed by Pieter Noordhuis
parent 9005896c0a
commit bc000c1db0
2 changed files with 42 additions and 22 deletions

48
redis.c
View File

@ -5922,6 +5922,7 @@ static int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
#define REDIS_AGGR_SUM 1 #define REDIS_AGGR_SUM 1
#define REDIS_AGGR_MIN 2 #define REDIS_AGGR_MIN 2
#define REDIS_AGGR_MAX 3 #define REDIS_AGGR_MAX 3
#define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e))
inline static void zunionInterAggregate(double *target, double val, int aggregate) { inline static void zunionInterAggregate(double *target, double val, int aggregate) {
if (aggregate == REDIS_AGGR_SUM) { if (aggregate == REDIS_AGGR_SUM) {
@ -5937,7 +5938,7 @@ inline static void zunionInterAggregate(double *target, double val, int aggregat
} }
static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) { static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
int i, j, zsetnum; int i, j, setnum;
int aggregate = REDIS_AGGR_SUM; int aggregate = REDIS_AGGR_SUM;
zsetopsrc *src; zsetopsrc *src;
robj *dstobj; robj *dstobj;
@ -5945,32 +5946,35 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
dictIterator *di; dictIterator *di;
dictEntry *de; dictEntry *de;
/* expect zsetnum input keys to be given */ /* expect setnum input keys to be given */
zsetnum = atoi(c->argv[2]->ptr); setnum = atoi(c->argv[2]->ptr);
if (zsetnum < 1) { if (setnum < 1) {
addReplySds(c,sdsnew("-ERR at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE\r\n")); addReplySds(c,sdsnew("-ERR at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE\r\n"));
return; return;
} }
/* test if the expected number of keys would overflow */ /* test if the expected number of keys would overflow */
if (3+zsetnum > c->argc) { if (3+setnum > c->argc) {
addReply(c,shared.syntaxerr); addReply(c,shared.syntaxerr);
return; return;
} }
/* read keys to be used for input */ /* read keys to be used for input */
src = zmalloc(sizeof(zsetopsrc) * zsetnum); src = zmalloc(sizeof(zsetopsrc) * setnum);
for (i = 0, j = 3; i < zsetnum; i++, j++) { for (i = 0, j = 3; i < setnum; i++, j++) {
robj *zsetobj = lookupKeyWrite(c->db,c->argv[j]); robj *obj = lookupKeyWrite(c->db,c->argv[j]);
if (!zsetobj) { if (!obj) {
src[i].dict = NULL; src[i].dict = NULL;
} else { } else {
if (zsetobj->type != REDIS_ZSET) { if (obj->type == REDIS_ZSET) {
src[i].dict = ((zset*)obj->ptr)->dict;
} else if (obj->type == REDIS_SET) {
src[i].dict = (obj->ptr);
} else {
zfree(src); zfree(src);
addReply(c,shared.wrongtypeerr); addReply(c,shared.wrongtypeerr);
return; return;
} }
src[i].dict = ((zset*)zsetobj->ptr)->dict;
} }
/* default all weights to 1 */ /* default all weights to 1 */
@ -5982,9 +5986,9 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
int remaining = c->argc - j; int remaining = c->argc - j;
while (remaining) { while (remaining) {
if (remaining >= (zsetnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) { if (remaining >= (setnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) {
j++; remaining--; j++; remaining--;
for (i = 0; i < zsetnum; i++, j++, remaining--) { for (i = 0; i < setnum; i++, j++, remaining--) {
if (getDoubleFromObjectOrReply(c, c->argv[j], &src[i].weight, NULL) != REDIS_OK) if (getDoubleFromObjectOrReply(c, c->argv[j], &src[i].weight, NULL) != REDIS_OK)
return; return;
} }
@ -6012,7 +6016,7 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
/* sort sets from the smallest to largest, this will improve our /* sort sets from the smallest to largest, this will improve our
* algorithm's performance */ * algorithm's performance */
qsort(src,zsetnum,sizeof(zsetopsrc), qsortCompareZsetopsrcByCardinality); qsort(src,setnum,sizeof(zsetopsrc),qsortCompareZsetopsrcByCardinality);
dstobj = createZsetObject(); dstobj = createZsetObject();
dstzset = dstobj->ptr; dstzset = dstobj->ptr;
@ -6025,12 +6029,12 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
di = dictGetIterator(src[0].dict); di = dictGetIterator(src[0].dict);
while((de = dictNext(di)) != NULL) { while((de = dictNext(di)) != NULL) {
double *score = zmalloc(sizeof(double)), value; double *score = zmalloc(sizeof(double)), value;
*score = src[0].weight * (*(double*)dictGetEntryVal(de)); *score = src[0].weight * zunionInterDictValue(de);
for (j = 1; j < zsetnum; j++) { for (j = 1; j < setnum; j++) {
dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de)); dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
if (other) { if (other) {
value = src[j].weight * (*(double*)dictGetEntryVal(other)); value = src[j].weight * zunionInterDictValue(other);
zunionInterAggregate(score, value, aggregate); zunionInterAggregate(score, value, aggregate);
} else { } else {
break; break;
@ -6038,7 +6042,7 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
} }
/* skip entry when not present in every source dict */ /* skip entry when not present in every source dict */
if (j != zsetnum) { if (j != setnum) {
zfree(score); zfree(score);
} else { } else {
robj *o = dictGetEntryKey(de); robj *o = dictGetEntryKey(de);
@ -6051,7 +6055,7 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
dictReleaseIterator(di); dictReleaseIterator(di);
} }
} else if (op == REDIS_OP_UNION) { } else if (op == REDIS_OP_UNION) {
for (i = 0; i < zsetnum; i++) { for (i = 0; i < setnum; i++) {
if (!src[i].dict) continue; if (!src[i].dict) continue;
di = dictGetIterator(src[i].dict); di = dictGetIterator(src[i].dict);
@ -6060,14 +6064,14 @@ static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) continue; if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) continue;
double *score = zmalloc(sizeof(double)), value; double *score = zmalloc(sizeof(double)), value;
*score = src[i].weight * (*(double*)dictGetEntryVal(de)); *score = src[i].weight * zunionInterDictValue(de);
/* because the zsets are sorted by size, its only possible /* because the zsets are sorted by size, its only possible
* for sets at larger indices to hold this entry */ * for sets at larger indices to hold this entry */
for (j = (i+1); j < zsetnum; j++) { for (j = (i+1); j < setnum; j++) {
dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de)); dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
if (other) { if (other) {
value = src[j].weight * (*(double*)dictGetEntryVal(other)); value = src[j].weight * zunionInterDictValue(other);
zunionInterAggregate(score, value, aggregate); zunionInterAggregate(score, value, aggregate);
} }
} }

View File

@ -316,6 +316,14 @@ start_server default.conf {} {
list [r zunionstore zsetc 2 zseta zsetb weights 2 3] [r zrange zsetc 0 -1 withscores] list [r zunionstore zsetc 2 zseta zsetb weights 2 3] [r zrange zsetc 0 -1 withscores]
} {4 {a 2 b 7 d 9 c 12}} } {4 {a 2 b 7 d 9 c 12}}
test {ZUNIONSTORE with a regular set and weights} {
r del seta
r sadd seta a
r sadd seta b
r sadd seta c
list [r zunionstore zsetc 2 seta zsetb weights 2 3] [r zrange zsetc 0 -1 withscores]
} {4 {a 2 b 5 c 8 d 9}}
test {ZUNIONSTORE with AGGREGATE MIN} { test {ZUNIONSTORE with AGGREGATE MIN} {
list [r zunionstore zsetc 2 zseta zsetb aggregate min] [r zrange zsetc 0 -1 withscores] list [r zunionstore zsetc 2 zseta zsetb aggregate min] [r zrange zsetc 0 -1 withscores]
} {4 {a 1 b 1 c 2 d 3}} } {4 {a 1 b 1 c 2 d 3}}
@ -332,6 +340,14 @@ start_server default.conf {} {
list [r zinterstore zsetc 2 zseta zsetb weights 2 3] [r zrange zsetc 0 -1 withscores] list [r zinterstore zsetc 2 zseta zsetb weights 2 3] [r zrange zsetc 0 -1 withscores]
} {2 {b 7 c 12}} } {2 {b 7 c 12}}
test {ZINTERSTORE with a regular set and weights} {
r del seta
r sadd seta a
r sadd seta b
r sadd seta c
list [r zinterstore zsetc 2 seta zsetb weights 2 3] [r zrange zsetc 0 -1 withscores]
} {2 {b 5 c 8}}
test {ZINTERSTORE with AGGREGATE MIN} { test {ZINTERSTORE with AGGREGATE MIN} {
list [r zinterstore zsetc 2 zseta zsetb aggregate min] [r zrange zsetc 0 -1 withscores] list [r zinterstore zsetc 2 zseta zsetb aggregate min] [r zrange zsetc 0 -1 withscores]
} {2 {b 1 c 2}} } {2 {b 1 c 2}}