+static void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
+ int i, j, k, zsetnum;
+ dict **srcdicts;
+ double *weights;
+ 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 */
+ srcdicts = zmalloc(sizeof(dict*) * zsetnum);
+ weights = zmalloc(sizeof(double) * zsetnum);
+ for (i = 0, j = 3; i < zsetnum; i++, j++) {
+ robj *zsetobj = lookupKeyWrite(c->db,c->argv[j]);
+ if (!zsetobj) {
+ srcdicts[i] = NULL;
+ } else {
+ if (zsetobj->type != REDIS_ZSET) {
+ zfree(srcdicts);
+ zfree(weights);
+ addReply(c,shared.wrongtypeerr);
+ return;
+ }
+ srcdicts[i] = ((zset*)zsetobj->ptr)->dict;
+ }
+
+ /* default all weights to 1 */
+ weights[i] = 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(srcdicts);
+ zfree(weights);
+ addReplySds(c,sdsnew("-ERR not enough weights for ZUNION/ZINTER\r\n"));
+ return;
+ }
+ for (i = 0; i < zsetnum; i++, j++, remaining--) {
+ weights[i] = strtod(c->argv[j]->ptr, NULL);
+ }
+ } else {
+ zfree(srcdicts);
+ zfree(weights);
+ addReply(c,shared.syntaxerr);
+ return;
+ }
+ }
+ }
+
+ dstobj = createZsetObject();
+ dstzset = dstobj->ptr;
+
+ if (op == REDIS_OP_INTER) {
+ /* store index of smallest zset in variable j */
+ for (i = 0, j = 0; i < zsetnum; i++) {
+ if (!srcdicts[i] || dictSize(srcdicts[i]) == 0) {
+ break;
+ }
+ if (dictSize(srcdicts[i]) < dictSize(srcdicts[j])) {
+ j = i;
+ }
+ }
+ /* 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) {
+ double *score = zmalloc(sizeof(double));
+ *score = 0.0;
+
+ for (k = 0; k < zsetnum; k++) {
+ dictEntry *other = (k == j) ? de : dictFind(srcdicts[k],dictGetEntryKey(de));
+ if (other) {
+ *score = *score + weights[k] * (*(double*)dictGetEntryVal(other));
+ } else {
+ break;
+ }
+ }
+
+ /* skip entry when not present in every source dict */
+ if (k != 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 (!srcdicts[i]) continue;
+
+ di = dictGetIterator(srcdicts[i]);
+ 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 (!srcdicts[j]) continue;
+
+ dictEntry *other = (i == j) ? de : dictFind(srcdicts[j],dictGetEntryKey(de));
+ if (other) {
+ *score = *score + weights[j] * (*(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(srcdicts);
+ zfree(weights);
+}
+
+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);
+}
+
projects / redis.git / blobdiff