]> git.saurik.com Git - redis.git/blame - src/t_zset.c
Fix ZUNIONSTORE/ZINTERSTORE to never store a NaN score.
[redis.git] / src / t_zset.c
CommitLineData
e2641e09 1#include "redis.h"
2
3#include <math.h>
4
5/*-----------------------------------------------------------------------------
6 * Sorted set API
7 *----------------------------------------------------------------------------*/
8
9/* ZSETs are ordered sets using two data structures to hold the same elements
10 * in order to get O(log(N)) INSERT and REMOVE operations into a sorted
11 * data structure.
12 *
13 * The elements are added to an hash table mapping Redis objects to scores.
14 * At the same time the elements are added to a skip list mapping scores
15 * to Redis objects (so objects are sorted by scores in this "view"). */
16
17/* This skiplist implementation is almost a C translation of the original
18 * algorithm described by William Pugh in "Skip Lists: A Probabilistic
19 * Alternative to Balanced Trees", modified in three ways:
20 * a) this implementation allows for repeated values.
21 * b) the comparison is not just by key (our 'score') but by satellite data.
22 * c) there is a back pointer, so it's a doubly linked list with the back
23 * pointers being only at "level 1". This allows to traverse the list
24 * from tail to head, useful for ZREVRANGE. */
25
26zskiplistNode *zslCreateNode(int level, double score, robj *obj) {
27 zskiplistNode *zn = zmalloc(sizeof(*zn));
28
29 zn->forward = zmalloc(sizeof(zskiplistNode*) * level);
30 if (level > 1)
31 zn->span = zmalloc(sizeof(unsigned int) * (level - 1));
32 else
33 zn->span = NULL;
34 zn->score = score;
35 zn->obj = obj;
36 return zn;
37}
38
39zskiplist *zslCreate(void) {
40 int j;
41 zskiplist *zsl;
42
43 zsl = zmalloc(sizeof(*zsl));
44 zsl->level = 1;
45 zsl->length = 0;
46 zsl->header = zslCreateNode(ZSKIPLIST_MAXLEVEL,0,NULL);
47 for (j = 0; j < ZSKIPLIST_MAXLEVEL; j++) {
48 zsl->header->forward[j] = NULL;
49
50 /* span has space for ZSKIPLIST_MAXLEVEL-1 elements */
51 if (j < ZSKIPLIST_MAXLEVEL-1)
52 zsl->header->span[j] = 0;
53 }
54 zsl->header->backward = NULL;
55 zsl->tail = NULL;
56 return zsl;
57}
58
59void zslFreeNode(zskiplistNode *node) {
60 decrRefCount(node->obj);
61 zfree(node->forward);
62 zfree(node->span);
63 zfree(node);
64}
65
66void zslFree(zskiplist *zsl) {
67 zskiplistNode *node = zsl->header->forward[0], *next;
68
69 zfree(zsl->header->forward);
70 zfree(zsl->header->span);
71 zfree(zsl->header);
72 while(node) {
73 next = node->forward[0];
74 zslFreeNode(node);
75 node = next;
76 }
77 zfree(zsl);
78}
79
80int zslRandomLevel(void) {
81 int level = 1;
82 while ((random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF))
83 level += 1;
84 return (level<ZSKIPLIST_MAXLEVEL) ? level : ZSKIPLIST_MAXLEVEL;
85}
86
87void zslInsert(zskiplist *zsl, double score, robj *obj) {
88 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
89 unsigned int rank[ZSKIPLIST_MAXLEVEL];
90 int i, level;
91
92 x = zsl->header;
93 for (i = zsl->level-1; i >= 0; i--) {
94 /* store rank that is crossed to reach the insert position */
95 rank[i] = i == (zsl->level-1) ? 0 : rank[i+1];
96
97 while (x->forward[i] &&
98 (x->forward[i]->score < score ||
99 (x->forward[i]->score == score &&
100 compareStringObjects(x->forward[i]->obj,obj) < 0))) {
101 rank[i] += i > 0 ? x->span[i-1] : 1;
102 x = x->forward[i];
103 }
104 update[i] = x;
105 }
106 /* we assume the key is not already inside, since we allow duplicated
107 * scores, and the re-insertion of score and redis object should never
108 * happpen since the caller of zslInsert() should test in the hash table
109 * if the element is already inside or not. */
110 level = zslRandomLevel();
111 if (level > zsl->level) {
112 for (i = zsl->level; i < level; i++) {
113 rank[i] = 0;
114 update[i] = zsl->header;
115 update[i]->span[i-1] = zsl->length;
116 }
117 zsl->level = level;
118 }
119 x = zslCreateNode(level,score,obj);
120 for (i = 0; i < level; i++) {
121 x->forward[i] = update[i]->forward[i];
122 update[i]->forward[i] = x;
123
124 /* update span covered by update[i] as x is inserted here */
125 if (i > 0) {
126 x->span[i-1] = update[i]->span[i-1] - (rank[0] - rank[i]);
127 update[i]->span[i-1] = (rank[0] - rank[i]) + 1;
128 }
129 }
130
131 /* increment span for untouched levels */
132 for (i = level; i < zsl->level; i++) {
133 update[i]->span[i-1]++;
134 }
135
136 x->backward = (update[0] == zsl->header) ? NULL : update[0];
137 if (x->forward[0])
138 x->forward[0]->backward = x;
139 else
140 zsl->tail = x;
141 zsl->length++;
142}
143
144/* Internal function used by zslDelete, zslDeleteByScore and zslDeleteByRank */
145void zslDeleteNode(zskiplist *zsl, zskiplistNode *x, zskiplistNode **update) {
146 int i;
147 for (i = 0; i < zsl->level; i++) {
148 if (update[i]->forward[i] == x) {
149 if (i > 0) {
150 update[i]->span[i-1] += x->span[i-1] - 1;
151 }
152 update[i]->forward[i] = x->forward[i];
153 } else {
154 /* invariant: i > 0, because update[0]->forward[0]
155 * is always equal to x */
156 update[i]->span[i-1] -= 1;
157 }
158 }
159 if (x->forward[0]) {
160 x->forward[0]->backward = x->backward;
161 } else {
162 zsl->tail = x->backward;
163 }
164 while(zsl->level > 1 && zsl->header->forward[zsl->level-1] == NULL)
165 zsl->level--;
166 zsl->length--;
167}
168
169/* Delete an element with matching score/object from the skiplist. */
170int zslDelete(zskiplist *zsl, double score, robj *obj) {
171 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
172 int i;
173
174 x = zsl->header;
175 for (i = zsl->level-1; i >= 0; i--) {
176 while (x->forward[i] &&
177 (x->forward[i]->score < score ||
178 (x->forward[i]->score == score &&
179 compareStringObjects(x->forward[i]->obj,obj) < 0)))
180 x = x->forward[i];
181 update[i] = x;
182 }
183 /* We may have multiple elements with the same score, what we need
184 * is to find the element with both the right score and object. */
185 x = x->forward[0];
186 if (x && score == x->score && equalStringObjects(x->obj,obj)) {
187 zslDeleteNode(zsl, x, update);
188 zslFreeNode(x);
189 return 1;
190 } else {
191 return 0; /* not found */
192 }
193 return 0; /* not found */
194}
195
196/* Delete all the elements with score between min and max from the skiplist.
197 * Min and mx are inclusive, so a score >= min || score <= max is deleted.
198 * Note that this function takes the reference to the hash table view of the
199 * sorted set, in order to remove the elements from the hash table too. */
200unsigned long zslDeleteRangeByScore(zskiplist *zsl, double min, double max, dict *dict) {
201 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
202 unsigned long removed = 0;
203 int i;
204
205 x = zsl->header;
206 for (i = zsl->level-1; i >= 0; i--) {
207 while (x->forward[i] && x->forward[i]->score < min)
208 x = x->forward[i];
209 update[i] = x;
210 }
211 /* We may have multiple elements with the same score, what we need
212 * is to find the element with both the right score and object. */
213 x = x->forward[0];
214 while (x && x->score <= max) {
215 zskiplistNode *next = x->forward[0];
216 zslDeleteNode(zsl, x, update);
217 dictDelete(dict,x->obj);
218 zslFreeNode(x);
219 removed++;
220 x = next;
221 }
222 return removed; /* not found */
223}
224
225/* Delete all the elements with rank between start and end from the skiplist.
226 * Start and end are inclusive. Note that start and end need to be 1-based */
227unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned int end, dict *dict) {
228 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
229 unsigned long traversed = 0, removed = 0;
230 int i;
231
232 x = zsl->header;
233 for (i = zsl->level-1; i >= 0; i--) {
234 while (x->forward[i] && (traversed + (i > 0 ? x->span[i-1] : 1)) < start) {
235 traversed += i > 0 ? x->span[i-1] : 1;
236 x = x->forward[i];
237 }
238 update[i] = x;
239 }
240
241 traversed++;
242 x = x->forward[0];
243 while (x && traversed <= end) {
244 zskiplistNode *next = x->forward[0];
245 zslDeleteNode(zsl, x, update);
246 dictDelete(dict,x->obj);
247 zslFreeNode(x);
248 removed++;
249 traversed++;
250 x = next;
251 }
252 return removed;
253}
254
255/* Find the first node having a score equal or greater than the specified one.
256 * Returns NULL if there is no match. */
257zskiplistNode *zslFirstWithScore(zskiplist *zsl, double score) {
258 zskiplistNode *x;
259 int i;
260
261 x = zsl->header;
262 for (i = zsl->level-1; i >= 0; i--) {
263 while (x->forward[i] && x->forward[i]->score < score)
264 x = x->forward[i];
265 }
266 /* We may have multiple elements with the same score, what we need
267 * is to find the element with both the right score and object. */
268 return x->forward[0];
269}
270
271/* Find the rank for an element by both score and key.
272 * Returns 0 when the element cannot be found, rank otherwise.
273 * Note that the rank is 1-based due to the span of zsl->header to the
274 * first element. */
275unsigned long zslistTypeGetRank(zskiplist *zsl, double score, robj *o) {
276 zskiplistNode *x;
277 unsigned long rank = 0;
278 int i;
279
280 x = zsl->header;
281 for (i = zsl->level-1; i >= 0; i--) {
282 while (x->forward[i] &&
283 (x->forward[i]->score < score ||
284 (x->forward[i]->score == score &&
285 compareStringObjects(x->forward[i]->obj,o) <= 0))) {
286 rank += i > 0 ? x->span[i-1] : 1;
287 x = x->forward[i];
288 }
289
290 /* x might be equal to zsl->header, so test if obj is non-NULL */
291 if (x->obj && equalStringObjects(x->obj,o)) {
292 return rank;
293 }
294 }
295 return 0;
296}
297
298/* Finds an element by its rank. The rank argument needs to be 1-based. */
299zskiplistNode* zslistTypeGetElementByRank(zskiplist *zsl, unsigned long rank) {
300 zskiplistNode *x;
301 unsigned long traversed = 0;
302 int i;
303
304 x = zsl->header;
305 for (i = zsl->level-1; i >= 0; i--) {
306 while (x->forward[i] && (traversed + (i>0 ? x->span[i-1] : 1)) <= rank)
307 {
308 traversed += i > 0 ? x->span[i-1] : 1;
309 x = x->forward[i];
310 }
311 if (traversed == rank) {
312 return x;
313 }
314 }
315 return NULL;
316}
317
318/*-----------------------------------------------------------------------------
319 * Sorted set commands
320 *----------------------------------------------------------------------------*/
321
322/* This generic command implements both ZADD and ZINCRBY.
323 * scoreval is the score if the operation is a ZADD (doincrement == 0) or
324 * the increment if the operation is a ZINCRBY (doincrement == 1). */
325void zaddGenericCommand(redisClient *c, robj *key, robj *ele, double scoreval, int doincrement) {
326 robj *zsetobj;
327 zset *zs;
328 double *score;
329
330 if (isnan(scoreval)) {
331 addReplySds(c,sdsnew("-ERR provide score is Not A Number (nan)\r\n"));
332 return;
333 }
334
335 zsetobj = lookupKeyWrite(c->db,key);
336 if (zsetobj == NULL) {
337 zsetobj = createZsetObject();
338 dbAdd(c->db,key,zsetobj);
339 } else {
340 if (zsetobj->type != REDIS_ZSET) {
341 addReply(c,shared.wrongtypeerr);
342 return;
343 }
344 }
345 zs = zsetobj->ptr;
346
347 /* Ok now since we implement both ZADD and ZINCRBY here the code
348 * needs to handle the two different conditions. It's all about setting
349 * '*score', that is, the new score to set, to the right value. */
350 score = zmalloc(sizeof(double));
351 if (doincrement) {
352 dictEntry *de;
353
354 /* Read the old score. If the element was not present starts from 0 */
355 de = dictFind(zs->dict,ele);
356 if (de) {
357 double *oldscore = dictGetEntryVal(de);
358 *score = *oldscore + scoreval;
359 } else {
360 *score = scoreval;
361 }
362 if (isnan(*score)) {
363 addReplySds(c,
364 sdsnew("-ERR resulting score is Not A Number (nan)\r\n"));
365 zfree(score);
366 /* Note that we don't need to check if the zset may be empty and
367 * should be removed here, as we can only obtain Nan as score if
368 * there was already an element in the sorted set. */
369 return;
370 }
371 } else {
372 *score = scoreval;
373 }
374
375 /* What follows is a simple remove and re-insert operation that is common
376 * to both ZADD and ZINCRBY... */
377 if (dictAdd(zs->dict,ele,score) == DICT_OK) {
378 /* case 1: New element */
379 incrRefCount(ele); /* added to hash */
380 zslInsert(zs->zsl,*score,ele);
381 incrRefCount(ele); /* added to skiplist */
5b4bff9c 382 touchWatchedKey(c->db,c->argv[1]);
e2641e09 383 server.dirty++;
384 if (doincrement)
385 addReplyDouble(c,*score);
386 else
387 addReply(c,shared.cone);
388 } else {
389 dictEntry *de;
390 double *oldscore;
391
392 /* case 2: Score update operation */
393 de = dictFind(zs->dict,ele);
394 redisAssert(de != NULL);
395 oldscore = dictGetEntryVal(de);
396 if (*score != *oldscore) {
397 int deleted;
398
399 /* Remove and insert the element in the skip list with new score */
400 deleted = zslDelete(zs->zsl,*oldscore,ele);
401 redisAssert(deleted != 0);
402 zslInsert(zs->zsl,*score,ele);
403 incrRefCount(ele);
404 /* Update the score in the hash table */
405 dictReplace(zs->dict,ele,score);
5b4bff9c 406 touchWatchedKey(c->db,c->argv[1]);
e2641e09 407 server.dirty++;
408 } else {
409 zfree(score);
410 }
411 if (doincrement)
412 addReplyDouble(c,*score);
413 else
414 addReply(c,shared.czero);
415 }
416}
417
418void zaddCommand(redisClient *c) {
419 double scoreval;
420
421 if (getDoubleFromObjectOrReply(c, c->argv[2], &scoreval, NULL) != REDIS_OK) return;
422 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,0);
423}
424
425void zincrbyCommand(redisClient *c) {
426 double scoreval;
427
428 if (getDoubleFromObjectOrReply(c, c->argv[2], &scoreval, NULL) != REDIS_OK) return;
429 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,1);
430}
431
432void zremCommand(redisClient *c) {
433 robj *zsetobj;
434 zset *zs;
435 dictEntry *de;
436 double *oldscore;
437 int deleted;
438
439 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
440 checkType(c,zsetobj,REDIS_ZSET)) return;
441
442 zs = zsetobj->ptr;
443 de = dictFind(zs->dict,c->argv[2]);
444 if (de == NULL) {
445 addReply(c,shared.czero);
446 return;
447 }
448 /* Delete from the skiplist */
449 oldscore = dictGetEntryVal(de);
450 deleted = zslDelete(zs->zsl,*oldscore,c->argv[2]);
451 redisAssert(deleted != 0);
452
453 /* Delete from the hash table */
454 dictDelete(zs->dict,c->argv[2]);
455 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
456 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
5b4bff9c 457 touchWatchedKey(c->db,c->argv[1]);
e2641e09 458 server.dirty++;
459 addReply(c,shared.cone);
460}
461
462void zremrangebyscoreCommand(redisClient *c) {
463 double min;
464 double max;
465 long deleted;
466 robj *zsetobj;
467 zset *zs;
468
469 if ((getDoubleFromObjectOrReply(c, c->argv[2], &min, NULL) != REDIS_OK) ||
470 (getDoubleFromObjectOrReply(c, c->argv[3], &max, NULL) != REDIS_OK)) return;
471
472 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
473 checkType(c,zsetobj,REDIS_ZSET)) return;
474
475 zs = zsetobj->ptr;
476 deleted = zslDeleteRangeByScore(zs->zsl,min,max,zs->dict);
477 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
478 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
5b4bff9c 479 if (deleted) touchWatchedKey(c->db,c->argv[1]);
e2641e09 480 server.dirty += deleted;
481 addReplyLongLong(c,deleted);
482}
483
484void zremrangebyrankCommand(redisClient *c) {
485 long start;
486 long end;
487 int llen;
488 long deleted;
489 robj *zsetobj;
490 zset *zs;
491
492 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
493 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
494
495 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
496 checkType(c,zsetobj,REDIS_ZSET)) return;
497 zs = zsetobj->ptr;
498 llen = zs->zsl->length;
499
500 /* convert negative indexes */
501 if (start < 0) start = llen+start;
502 if (end < 0) end = llen+end;
503 if (start < 0) start = 0;
e2641e09 504
d0a4e24e
PN
505 /* Invariant: start >= 0, so this test will be true when end < 0.
506 * The range is empty when start > end or start >= length. */
e2641e09 507 if (start > end || start >= llen) {
508 addReply(c,shared.czero);
509 return;
510 }
511 if (end >= llen) end = llen-1;
512
513 /* increment start and end because zsl*Rank functions
514 * use 1-based rank */
515 deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
516 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
517 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
5b4bff9c 518 if (deleted) touchWatchedKey(c->db,c->argv[1]);
e2641e09 519 server.dirty += deleted;
520 addReplyLongLong(c, deleted);
521}
522
523typedef struct {
524 dict *dict;
525 double weight;
526} zsetopsrc;
527
528int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
529 zsetopsrc *d1 = (void*) s1, *d2 = (void*) s2;
530 unsigned long size1, size2;
531 size1 = d1->dict ? dictSize(d1->dict) : 0;
532 size2 = d2->dict ? dictSize(d2->dict) : 0;
533 return size1 - size2;
534}
535
536#define REDIS_AGGR_SUM 1
537#define REDIS_AGGR_MIN 2
538#define REDIS_AGGR_MAX 3
539#define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e))
540
541inline static void zunionInterAggregate(double *target, double val, int aggregate) {
542 if (aggregate == REDIS_AGGR_SUM) {
543 *target = *target + val;
d9e28bcf
PN
544 /* The result of adding two doubles is NaN when one variable
545 * is +inf and the other is -inf. When these numbers are added,
546 * we maintain the convention of the result being 0.0. */
547 if (isnan(*target)) *target = 0.0;
e2641e09 548 } else if (aggregate == REDIS_AGGR_MIN) {
549 *target = val < *target ? val : *target;
550 } else if (aggregate == REDIS_AGGR_MAX) {
551 *target = val > *target ? val : *target;
552 } else {
553 /* safety net */
554 redisPanic("Unknown ZUNION/INTER aggregate type");
555 }
556}
557
558void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
559 int i, j, setnum;
560 int aggregate = REDIS_AGGR_SUM;
561 zsetopsrc *src;
562 robj *dstobj;
563 zset *dstzset;
564 dictIterator *di;
565 dictEntry *de;
8c1420ff 566 int touched = 0;
e2641e09 567
568 /* expect setnum input keys to be given */
569 setnum = atoi(c->argv[2]->ptr);
570 if (setnum < 1) {
571 addReplySds(c,sdsnew("-ERR at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE\r\n"));
572 return;
573 }
574
575 /* test if the expected number of keys would overflow */
576 if (3+setnum > c->argc) {
577 addReply(c,shared.syntaxerr);
578 return;
579 }
580
581 /* read keys to be used for input */
582 src = zmalloc(sizeof(zsetopsrc) * setnum);
583 for (i = 0, j = 3; i < setnum; i++, j++) {
584 robj *obj = lookupKeyWrite(c->db,c->argv[j]);
585 if (!obj) {
586 src[i].dict = NULL;
587 } else {
588 if (obj->type == REDIS_ZSET) {
589 src[i].dict = ((zset*)obj->ptr)->dict;
590 } else if (obj->type == REDIS_SET) {
591 src[i].dict = (obj->ptr);
592 } else {
593 zfree(src);
594 addReply(c,shared.wrongtypeerr);
595 return;
596 }
597 }
598
599 /* default all weights to 1 */
600 src[i].weight = 1.0;
601 }
602
603 /* parse optional extra arguments */
604 if (j < c->argc) {
605 int remaining = c->argc - j;
606
607 while (remaining) {
608 if (remaining >= (setnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) {
609 j++; remaining--;
610 for (i = 0; i < setnum; i++, j++, remaining--) {
611 if (getDoubleFromObjectOrReply(c, c->argv[j], &src[i].weight, NULL) != REDIS_OK)
612 return;
613 }
614 } else if (remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
615 j++; remaining--;
616 if (!strcasecmp(c->argv[j]->ptr,"sum")) {
617 aggregate = REDIS_AGGR_SUM;
618 } else if (!strcasecmp(c->argv[j]->ptr,"min")) {
619 aggregate = REDIS_AGGR_MIN;
620 } else if (!strcasecmp(c->argv[j]->ptr,"max")) {
621 aggregate = REDIS_AGGR_MAX;
622 } else {
623 zfree(src);
624 addReply(c,shared.syntaxerr);
625 return;
626 }
627 j++; remaining--;
628 } else {
629 zfree(src);
630 addReply(c,shared.syntaxerr);
631 return;
632 }
633 }
634 }
635
636 /* sort sets from the smallest to largest, this will improve our
637 * algorithm's performance */
638 qsort(src,setnum,sizeof(zsetopsrc),qsortCompareZsetopsrcByCardinality);
639
640 dstobj = createZsetObject();
641 dstzset = dstobj->ptr;
642
643 if (op == REDIS_OP_INTER) {
644 /* skip going over all entries if the smallest zset is NULL or empty */
645 if (src[0].dict && dictSize(src[0].dict) > 0) {
646 /* precondition: as src[0].dict is non-empty and the zsets are ordered
647 * from small to large, all src[i > 0].dict are non-empty too */
648 di = dictGetIterator(src[0].dict);
649 while((de = dictNext(di)) != NULL) {
650 double *score = zmalloc(sizeof(double)), value;
651 *score = src[0].weight * zunionInterDictValue(de);
652
653 for (j = 1; j < setnum; j++) {
654 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
655 if (other) {
656 value = src[j].weight * zunionInterDictValue(other);
657 zunionInterAggregate(score, value, aggregate);
658 } else {
659 break;
660 }
661 }
662
663 /* skip entry when not present in every source dict */
664 if (j != setnum) {
665 zfree(score);
666 } else {
667 robj *o = dictGetEntryKey(de);
668 dictAdd(dstzset->dict,o,score);
669 incrRefCount(o); /* added to dictionary */
670 zslInsert(dstzset->zsl,*score,o);
671 incrRefCount(o); /* added to skiplist */
672 }
673 }
674 dictReleaseIterator(di);
675 }
676 } else if (op == REDIS_OP_UNION) {
677 for (i = 0; i < setnum; i++) {
678 if (!src[i].dict) continue;
679
680 di = dictGetIterator(src[i].dict);
681 while((de = dictNext(di)) != NULL) {
682 /* skip key when already processed */
683 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) continue;
684
685 double *score = zmalloc(sizeof(double)), value;
686 *score = src[i].weight * zunionInterDictValue(de);
687
688 /* because the zsets are sorted by size, its only possible
689 * for sets at larger indices to hold this entry */
690 for (j = (i+1); j < setnum; j++) {
691 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
692 if (other) {
693 value = src[j].weight * zunionInterDictValue(other);
694 zunionInterAggregate(score, value, aggregate);
695 }
696 }
697
698 robj *o = dictGetEntryKey(de);
699 dictAdd(dstzset->dict,o,score);
700 incrRefCount(o); /* added to dictionary */
701 zslInsert(dstzset->zsl,*score,o);
702 incrRefCount(o); /* added to skiplist */
703 }
704 dictReleaseIterator(di);
705 }
706 } else {
707 /* unknown operator */
708 redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION);
709 }
710
8c1420ff 711 if (dbDelete(c->db,dstkey)) {
712 touchWatchedKey(c->db,dstkey);
713 touched = 1;
714 server.dirty++;
715 }
e2641e09 716 if (dstzset->zsl->length) {
717 dbAdd(c->db,dstkey,dstobj);
718 addReplyLongLong(c, dstzset->zsl->length);
8c1420ff 719 if (!touched) touchWatchedKey(c->db,dstkey);
cbf7e107 720 server.dirty++;
e2641e09 721 } else {
722 decrRefCount(dstobj);
723 addReply(c, shared.czero);
724 }
725 zfree(src);
726}
727
728void zunionstoreCommand(redisClient *c) {
729 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION);
730}
731
732void zinterstoreCommand(redisClient *c) {
733 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
734}
735
736void zrangeGenericCommand(redisClient *c, int reverse) {
737 robj *o;
738 long start;
739 long end;
740 int withscores = 0;
741 int llen;
742 int rangelen, j;
743 zset *zsetobj;
744 zskiplist *zsl;
745 zskiplistNode *ln;
746 robj *ele;
747
748 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
749 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
750
751 if (c->argc == 5 && !strcasecmp(c->argv[4]->ptr,"withscores")) {
752 withscores = 1;
753 } else if (c->argc >= 5) {
754 addReply(c,shared.syntaxerr);
755 return;
756 }
757
758 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) == NULL
759 || checkType(c,o,REDIS_ZSET)) return;
760 zsetobj = o->ptr;
761 zsl = zsetobj->zsl;
762 llen = zsl->length;
763
764 /* convert negative indexes */
765 if (start < 0) start = llen+start;
766 if (end < 0) end = llen+end;
767 if (start < 0) start = 0;
e2641e09 768
d0a4e24e
PN
769 /* Invariant: start >= 0, so this test will be true when end < 0.
770 * The range is empty when start > end or start >= length. */
e2641e09 771 if (start > end || start >= llen) {
e2641e09 772 addReply(c,shared.emptymultibulk);
773 return;
774 }
775 if (end >= llen) end = llen-1;
776 rangelen = (end-start)+1;
777
778 /* check if starting point is trivial, before searching
779 * the element in log(N) time */
780 if (reverse) {
781 ln = start == 0 ? zsl->tail : zslistTypeGetElementByRank(zsl, llen-start);
782 } else {
783 ln = start == 0 ?
784 zsl->header->forward[0] : zslistTypeGetElementByRank(zsl, start+1);
785 }
786
787 /* Return the result in form of a multi-bulk reply */
788 addReplySds(c,sdscatprintf(sdsempty(),"*%d\r\n",
789 withscores ? (rangelen*2) : rangelen));
790 for (j = 0; j < rangelen; j++) {
791 ele = ln->obj;
792 addReplyBulk(c,ele);
793 if (withscores)
794 addReplyDouble(c,ln->score);
795 ln = reverse ? ln->backward : ln->forward[0];
796 }
797}
798
799void zrangeCommand(redisClient *c) {
800 zrangeGenericCommand(c,0);
801}
802
803void zrevrangeCommand(redisClient *c) {
804 zrangeGenericCommand(c,1);
805}
806
807/* This command implements both ZRANGEBYSCORE and ZCOUNT.
808 * If justcount is non-zero, just the count is returned. */
809void genericZrangebyscoreCommand(redisClient *c, int justcount) {
810 robj *o;
811 double min, max;
812 int minex = 0, maxex = 0; /* are min or max exclusive? */
813 int offset = 0, limit = -1;
814 int withscores = 0;
815 int badsyntax = 0;
816
817 /* Parse the min-max interval. If one of the values is prefixed
818 * by the "(" character, it's considered "open". For instance
819 * ZRANGEBYSCORE zset (1.5 (2.5 will match min < x < max
820 * ZRANGEBYSCORE zset 1.5 2.5 will instead match min <= x <= max */
821 if (((char*)c->argv[2]->ptr)[0] == '(') {
822 min = strtod((char*)c->argv[2]->ptr+1,NULL);
823 minex = 1;
824 } else {
825 min = strtod(c->argv[2]->ptr,NULL);
826 }
827 if (((char*)c->argv[3]->ptr)[0] == '(') {
828 max = strtod((char*)c->argv[3]->ptr+1,NULL);
829 maxex = 1;
830 } else {
831 max = strtod(c->argv[3]->ptr,NULL);
832 }
833
834 /* Parse "WITHSCORES": note that if the command was called with
835 * the name ZCOUNT then we are sure that c->argc == 4, so we'll never
836 * enter the following paths to parse WITHSCORES and LIMIT. */
837 if (c->argc == 5 || c->argc == 8) {
838 if (strcasecmp(c->argv[c->argc-1]->ptr,"withscores") == 0)
839 withscores = 1;
840 else
841 badsyntax = 1;
842 }
843 if (c->argc != (4 + withscores) && c->argc != (7 + withscores))
844 badsyntax = 1;
845 if (badsyntax) {
846 addReplySds(c,
847 sdsnew("-ERR wrong number of arguments for ZRANGEBYSCORE\r\n"));
848 return;
849 }
850
851 /* Parse "LIMIT" */
852 if (c->argc == (7 + withscores) && strcasecmp(c->argv[4]->ptr,"limit")) {
853 addReply(c,shared.syntaxerr);
854 return;
855 } else if (c->argc == (7 + withscores)) {
856 offset = atoi(c->argv[5]->ptr);
857 limit = atoi(c->argv[6]->ptr);
858 if (offset < 0) offset = 0;
859 }
860
861 /* Ok, lookup the key and get the range */
862 o = lookupKeyRead(c->db,c->argv[1]);
863 if (o == NULL) {
864 addReply(c,justcount ? shared.czero : shared.emptymultibulk);
865 } else {
866 if (o->type != REDIS_ZSET) {
867 addReply(c,shared.wrongtypeerr);
868 } else {
869 zset *zsetobj = o->ptr;
870 zskiplist *zsl = zsetobj->zsl;
871 zskiplistNode *ln;
872 robj *ele, *lenobj = NULL;
873 unsigned long rangelen = 0;
874
875 /* Get the first node with the score >= min, or with
876 * score > min if 'minex' is true. */
877 ln = zslFirstWithScore(zsl,min);
878 while (minex && ln && ln->score == min) ln = ln->forward[0];
879
880 if (ln == NULL) {
881 /* No element matching the speciifed interval */
882 addReply(c,justcount ? shared.czero : shared.emptymultibulk);
883 return;
884 }
885
886 /* We don't know in advance how many matching elements there
887 * are in the list, so we push this object that will represent
888 * the multi-bulk length in the output buffer, and will "fix"
889 * it later */
890 if (!justcount) {
891 lenobj = createObject(REDIS_STRING,NULL);
892 addReply(c,lenobj);
893 decrRefCount(lenobj);
894 }
895
896 while(ln && (maxex ? (ln->score < max) : (ln->score <= max))) {
897 if (offset) {
898 offset--;
899 ln = ln->forward[0];
900 continue;
901 }
902 if (limit == 0) break;
903 if (!justcount) {
904 ele = ln->obj;
905 addReplyBulk(c,ele);
906 if (withscores)
907 addReplyDouble(c,ln->score);
908 }
909 ln = ln->forward[0];
910 rangelen++;
911 if (limit > 0) limit--;
912 }
913 if (justcount) {
914 addReplyLongLong(c,(long)rangelen);
915 } else {
916 lenobj->ptr = sdscatprintf(sdsempty(),"*%lu\r\n",
917 withscores ? (rangelen*2) : rangelen);
918 }
919 }
920 }
921}
922
923void zrangebyscoreCommand(redisClient *c) {
924 genericZrangebyscoreCommand(c,0);
925}
926
927void zcountCommand(redisClient *c) {
928 genericZrangebyscoreCommand(c,1);
929}
930
931void zcardCommand(redisClient *c) {
932 robj *o;
933 zset *zs;
934
935 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
936 checkType(c,o,REDIS_ZSET)) return;
937
938 zs = o->ptr;
939 addReplyUlong(c,zs->zsl->length);
940}
941
942void zscoreCommand(redisClient *c) {
943 robj *o;
944 zset *zs;
945 dictEntry *de;
946
947 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
948 checkType(c,o,REDIS_ZSET)) return;
949
950 zs = o->ptr;
951 de = dictFind(zs->dict,c->argv[2]);
952 if (!de) {
953 addReply(c,shared.nullbulk);
954 } else {
955 double *score = dictGetEntryVal(de);
956
957 addReplyDouble(c,*score);
958 }
959}
960
961void zrankGenericCommand(redisClient *c, int reverse) {
962 robj *o;
963 zset *zs;
964 zskiplist *zsl;
965 dictEntry *de;
966 unsigned long rank;
967 double *score;
968
969 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
970 checkType(c,o,REDIS_ZSET)) return;
971
972 zs = o->ptr;
973 zsl = zs->zsl;
974 de = dictFind(zs->dict,c->argv[2]);
975 if (!de) {
976 addReply(c,shared.nullbulk);
977 return;
978 }
979
980 score = dictGetEntryVal(de);
981 rank = zslistTypeGetRank(zsl, *score, c->argv[2]);
982 if (rank) {
983 if (reverse) {
984 addReplyLongLong(c, zsl->length - rank);
985 } else {
986 addReplyLongLong(c, rank-1);
987 }
988 } else {
989 addReply(c,shared.nullbulk);
990 }
991}
992
993void zrankCommand(redisClient *c) {
994 zrankGenericCommand(c, 0);
995}
996
997void zrevrankCommand(redisClient *c) {
998 zrankGenericCommand(c, 1);
999}