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