]> git.saurik.com Git - redis.git/blob - src/t_zset.c
Merge remote branch 'remotes/pietern/zrevrangebyscore'
[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 /* Struct to hold a inclusive/exclusive range spec. */
178 typedef struct {
179 double min, max;
180 int minex, maxex; /* are min or max exclusive? */
181 } zrangespec;
182
183 /* Delete all the elements with score between min and max from the skiplist.
184 * Min and mx are inclusive, so a score >= min || score <= max is deleted.
185 * Note that this function takes the reference to the hash table view of the
186 * sorted set, in order to remove the elements from the hash table too. */
187 unsigned long zslDeleteRangeByScore(zskiplist *zsl, zrangespec range, dict *dict) {
188 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
189 unsigned long removed = 0;
190 int i;
191
192 x = zsl->header;
193 for (i = zsl->level-1; i >= 0; i--) {
194 while (x->level[i].forward && (range.minex ?
195 x->level[i].forward->score <= range.min :
196 x->level[i].forward->score < range.min))
197 x = x->level[i].forward;
198 update[i] = x;
199 }
200
201 /* Current node is the last with score < or <= min. */
202 x = x->level[0].forward;
203
204 /* Delete nodes while in range. */
205 while (x && (range.maxex ? x->score < range.max : x->score <= range.max)) {
206 zskiplistNode *next = x->level[0].forward;
207 zslDeleteNode(zsl,x,update);
208 dictDelete(dict,x->obj);
209 zslFreeNode(x);
210 removed++;
211 x = next;
212 }
213 return removed;
214 }
215
216 /* Delete all the elements with rank between start and end from the skiplist.
217 * Start and end are inclusive. Note that start and end need to be 1-based */
218 unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned int end, dict *dict) {
219 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
220 unsigned long traversed = 0, removed = 0;
221 int i;
222
223 x = zsl->header;
224 for (i = zsl->level-1; i >= 0; i--) {
225 while (x->level[i].forward && (traversed + x->level[i].span) < start) {
226 traversed += x->level[i].span;
227 x = x->level[i].forward;
228 }
229 update[i] = x;
230 }
231
232 traversed++;
233 x = x->level[0].forward;
234 while (x && traversed <= end) {
235 zskiplistNode *next = x->level[0].forward;
236 zslDeleteNode(zsl,x,update);
237 dictDelete(dict,x->obj);
238 zslFreeNode(x);
239 removed++;
240 traversed++;
241 x = next;
242 }
243 return removed;
244 }
245
246 /* Find the first node having a score equal or greater than the specified one.
247 * Returns NULL if there is no match. */
248 zskiplistNode *zslFirstWithScore(zskiplist *zsl, double score) {
249 zskiplistNode *x;
250 int i;
251
252 x = zsl->header;
253 for (i = zsl->level-1; i >= 0; i--) {
254 while (x->level[i].forward && x->level[i].forward->score < score)
255 x = x->level[i].forward;
256 }
257 /* We may have multiple elements with the same score, what we need
258 * is to find the element with both the right score and object. */
259 return x->level[0].forward;
260 }
261
262 /* Find the rank for an element by both score and key.
263 * Returns 0 when the element cannot be found, rank otherwise.
264 * Note that the rank is 1-based due to the span of zsl->header to the
265 * first element. */
266 unsigned long zslistTypeGetRank(zskiplist *zsl, double score, robj *o) {
267 zskiplistNode *x;
268 unsigned long rank = 0;
269 int i;
270
271 x = zsl->header;
272 for (i = zsl->level-1; i >= 0; i--) {
273 while (x->level[i].forward &&
274 (x->level[i].forward->score < score ||
275 (x->level[i].forward->score == score &&
276 compareStringObjects(x->level[i].forward->obj,o) <= 0))) {
277 rank += x->level[i].span;
278 x = x->level[i].forward;
279 }
280
281 /* x might be equal to zsl->header, so test if obj is non-NULL */
282 if (x->obj && equalStringObjects(x->obj,o)) {
283 return rank;
284 }
285 }
286 return 0;
287 }
288
289 /* Finds an element by its rank. The rank argument needs to be 1-based. */
290 zskiplistNode* zslistTypeGetElementByRank(zskiplist *zsl, unsigned long rank) {
291 zskiplistNode *x;
292 unsigned long traversed = 0;
293 int i;
294
295 x = zsl->header;
296 for (i = zsl->level-1; i >= 0; i--) {
297 while (x->level[i].forward && (traversed + x->level[i].span) <= rank)
298 {
299 traversed += x->level[i].span;
300 x = x->level[i].forward;
301 }
302 if (traversed == rank) {
303 return x;
304 }
305 }
306 return NULL;
307 }
308
309 /* Populate the rangespec according to the objects min and max. */
310 static int zslParseRange(robj *min, robj *max, zrangespec *spec) {
311 char *eptr;
312 spec->minex = spec->maxex = 0;
313
314 /* Parse the min-max interval. If one of the values is prefixed
315 * by the "(" character, it's considered "open". For instance
316 * ZRANGEBYSCORE zset (1.5 (2.5 will match min < x < max
317 * ZRANGEBYSCORE zset 1.5 2.5 will instead match min <= x <= max */
318 if (min->encoding == REDIS_ENCODING_INT) {
319 spec->min = (long)min->ptr;
320 } else {
321 if (((char*)min->ptr)[0] == '(') {
322 spec->min = strtod((char*)min->ptr+1,&eptr);
323 if (eptr[0] != '\0' || isnan(spec->min)) return REDIS_ERR;
324 spec->minex = 1;
325 } else {
326 spec->min = strtod((char*)min->ptr,&eptr);
327 if (eptr[0] != '\0' || isnan(spec->min)) return REDIS_ERR;
328 }
329 }
330 if (max->encoding == REDIS_ENCODING_INT) {
331 spec->max = (long)max->ptr;
332 } else {
333 if (((char*)max->ptr)[0] == '(') {
334 spec->max = strtod((char*)max->ptr+1,&eptr);
335 if (eptr[0] != '\0' || isnan(spec->max)) return REDIS_ERR;
336 spec->maxex = 1;
337 } else {
338 spec->max = strtod((char*)max->ptr,&eptr);
339 if (eptr[0] != '\0' || isnan(spec->max)) return REDIS_ERR;
340 }
341 }
342
343 return REDIS_OK;
344 }
345
346
347 /*-----------------------------------------------------------------------------
348 * Sorted set commands
349 *----------------------------------------------------------------------------*/
350
351 /* This generic command implements both ZADD and ZINCRBY. */
352 void zaddGenericCommand(redisClient *c, robj *key, robj *ele, double score, int incr) {
353 robj *zsetobj;
354 zset *zs;
355 zskiplistNode *znode;
356
357 zsetobj = lookupKeyWrite(c->db,key);
358 if (zsetobj == NULL) {
359 zsetobj = createZsetObject();
360 dbAdd(c->db,key,zsetobj);
361 } else {
362 if (zsetobj->type != REDIS_ZSET) {
363 addReply(c,shared.wrongtypeerr);
364 return;
365 }
366 }
367 zs = zsetobj->ptr;
368
369 /* Since both ZADD and ZINCRBY are implemented here, we need to increment
370 * the score first by the current score if ZINCRBY is called. */
371 if (incr) {
372 /* Read the old score. If the element was not present starts from 0 */
373 dictEntry *de = dictFind(zs->dict,ele);
374 if (de != NULL)
375 score += *(double*)dictGetEntryVal(de);
376
377 if (isnan(score)) {
378 addReplyError(c,"resulting score is not a number (NaN)");
379 /* Note that we don't need to check if the zset may be empty and
380 * should be removed here, as we can only obtain Nan as score if
381 * there was already an element in the sorted set. */
382 return;
383 }
384 }
385
386 /* We need to remove and re-insert the element when it was already present
387 * in the dictionary, to update the skiplist. Note that we delay adding a
388 * pointer to the score because we want to reference the score in the
389 * skiplist node. */
390 if (dictAdd(zs->dict,ele,NULL) == DICT_OK) {
391 dictEntry *de;
392
393 /* New element */
394 incrRefCount(ele); /* added to hash */
395 znode = zslInsert(zs->zsl,score,ele);
396 incrRefCount(ele); /* added to skiplist */
397
398 /* Update the score in the dict entry */
399 de = dictFind(zs->dict,ele);
400 redisAssert(de != NULL);
401 dictGetEntryVal(de) = &znode->score;
402 touchWatchedKey(c->db,c->argv[1]);
403 server.dirty++;
404 if (incr)
405 addReplyDouble(c,score);
406 else
407 addReply(c,shared.cone);
408 } else {
409 dictEntry *de;
410 robj *curobj;
411 double *curscore;
412 int deleted;
413
414 /* Update score */
415 de = dictFind(zs->dict,ele);
416 redisAssert(de != NULL);
417 curobj = dictGetEntryKey(de);
418 curscore = dictGetEntryVal(de);
419
420 /* When the score is updated, reuse the existing string object to
421 * prevent extra alloc/dealloc of strings on ZINCRBY. */
422 if (score != *curscore) {
423 deleted = zslDelete(zs->zsl,*curscore,curobj);
424 redisAssert(deleted != 0);
425 znode = zslInsert(zs->zsl,score,curobj);
426 incrRefCount(curobj);
427
428 /* Update the score in the current dict entry */
429 dictGetEntryVal(de) = &znode->score;
430 touchWatchedKey(c->db,c->argv[1]);
431 server.dirty++;
432 }
433 if (incr)
434 addReplyDouble(c,score);
435 else
436 addReply(c,shared.czero);
437 }
438 }
439
440 void zaddCommand(redisClient *c) {
441 double scoreval;
442 if (getDoubleFromObjectOrReply(c,c->argv[2],&scoreval,NULL) != REDIS_OK) return;
443 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,0);
444 }
445
446 void zincrbyCommand(redisClient *c) {
447 double scoreval;
448 if (getDoubleFromObjectOrReply(c,c->argv[2],&scoreval,NULL) != REDIS_OK) return;
449 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,1);
450 }
451
452 void zremCommand(redisClient *c) {
453 robj *zsetobj;
454 zset *zs;
455 dictEntry *de;
456 double curscore;
457 int deleted;
458
459 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
460 checkType(c,zsetobj,REDIS_ZSET)) return;
461
462 zs = zsetobj->ptr;
463 de = dictFind(zs->dict,c->argv[2]);
464 if (de == NULL) {
465 addReply(c,shared.czero);
466 return;
467 }
468 /* Delete from the skiplist */
469 curscore = *(double*)dictGetEntryVal(de);
470 deleted = zslDelete(zs->zsl,curscore,c->argv[2]);
471 redisAssert(deleted != 0);
472
473 /* Delete from the hash table */
474 dictDelete(zs->dict,c->argv[2]);
475 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
476 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
477 touchWatchedKey(c->db,c->argv[1]);
478 server.dirty++;
479 addReply(c,shared.cone);
480 }
481
482 void zremrangebyscoreCommand(redisClient *c) {
483 zrangespec range;
484 long deleted;
485 robj *o;
486 zset *zs;
487
488 /* Parse the range arguments. */
489 if (zslParseRange(c->argv[2],c->argv[3],&range) != REDIS_OK) {
490 addReplyError(c,"min or max is not a double");
491 return;
492 }
493
494 if ((o = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
495 checkType(c,o,REDIS_ZSET)) return;
496
497 zs = o->ptr;
498 deleted = zslDeleteRangeByScore(zs->zsl,range,zs->dict);
499 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
500 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
501 if (deleted) touchWatchedKey(c->db,c->argv[1]);
502 server.dirty += deleted;
503 addReplyLongLong(c,deleted);
504 }
505
506 void zremrangebyrankCommand(redisClient *c) {
507 long start;
508 long end;
509 int llen;
510 long deleted;
511 robj *zsetobj;
512 zset *zs;
513
514 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
515 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
516
517 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
518 checkType(c,zsetobj,REDIS_ZSET)) return;
519 zs = zsetobj->ptr;
520 llen = zs->zsl->length;
521
522 /* convert negative indexes */
523 if (start < 0) start = llen+start;
524 if (end < 0) end = llen+end;
525 if (start < 0) start = 0;
526
527 /* Invariant: start >= 0, so this test will be true when end < 0.
528 * The range is empty when start > end or start >= length. */
529 if (start > end || start >= llen) {
530 addReply(c,shared.czero);
531 return;
532 }
533 if (end >= llen) end = llen-1;
534
535 /* increment start and end because zsl*Rank functions
536 * use 1-based rank */
537 deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
538 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
539 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
540 if (deleted) touchWatchedKey(c->db,c->argv[1]);
541 server.dirty += deleted;
542 addReplyLongLong(c, deleted);
543 }
544
545 typedef struct {
546 dict *dict;
547 double weight;
548 } zsetopsrc;
549
550 int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
551 zsetopsrc *d1 = (void*) s1, *d2 = (void*) s2;
552 unsigned long size1, size2;
553 size1 = d1->dict ? dictSize(d1->dict) : 0;
554 size2 = d2->dict ? dictSize(d2->dict) : 0;
555 return size1 - size2;
556 }
557
558 #define REDIS_AGGR_SUM 1
559 #define REDIS_AGGR_MIN 2
560 #define REDIS_AGGR_MAX 3
561 #define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e))
562
563 inline static void zunionInterAggregate(double *target, double val, int aggregate) {
564 if (aggregate == REDIS_AGGR_SUM) {
565 *target = *target + val;
566 /* The result of adding two doubles is NaN when one variable
567 * is +inf and the other is -inf. When these numbers are added,
568 * we maintain the convention of the result being 0.0. */
569 if (isnan(*target)) *target = 0.0;
570 } else if (aggregate == REDIS_AGGR_MIN) {
571 *target = val < *target ? val : *target;
572 } else if (aggregate == REDIS_AGGR_MAX) {
573 *target = val > *target ? val : *target;
574 } else {
575 /* safety net */
576 redisPanic("Unknown ZUNION/INTER aggregate type");
577 }
578 }
579
580 void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
581 int i, j, setnum;
582 int aggregate = REDIS_AGGR_SUM;
583 zsetopsrc *src;
584 robj *dstobj;
585 zset *dstzset;
586 zskiplistNode *znode;
587 dictIterator *di;
588 dictEntry *de;
589 int touched = 0;
590
591 /* expect setnum input keys to be given */
592 setnum = atoi(c->argv[2]->ptr);
593 if (setnum < 1) {
594 addReplyError(c,
595 "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE");
596 return;
597 }
598
599 /* test if the expected number of keys would overflow */
600 if (3+setnum > c->argc) {
601 addReply(c,shared.syntaxerr);
602 return;
603 }
604
605 /* read keys to be used for input */
606 src = zmalloc(sizeof(zsetopsrc) * setnum);
607 for (i = 0, j = 3; i < setnum; i++, j++) {
608 robj *obj = lookupKeyWrite(c->db,c->argv[j]);
609 if (!obj) {
610 src[i].dict = NULL;
611 } else {
612 if (obj->type == REDIS_ZSET) {
613 src[i].dict = ((zset*)obj->ptr)->dict;
614 } else if (obj->type == REDIS_SET) {
615 src[i].dict = (obj->ptr);
616 } else {
617 zfree(src);
618 addReply(c,shared.wrongtypeerr);
619 return;
620 }
621 }
622
623 /* default all weights to 1 */
624 src[i].weight = 1.0;
625 }
626
627 /* parse optional extra arguments */
628 if (j < c->argc) {
629 int remaining = c->argc - j;
630
631 while (remaining) {
632 if (remaining >= (setnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) {
633 j++; remaining--;
634 for (i = 0; i < setnum; i++, j++, remaining--) {
635 if (getDoubleFromObjectOrReply(c,c->argv[j],&src[i].weight,
636 "weight value is not a double") != REDIS_OK)
637 {
638 zfree(src);
639 return;
640 }
641 }
642 } else if (remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
643 j++; remaining--;
644 if (!strcasecmp(c->argv[j]->ptr,"sum")) {
645 aggregate = REDIS_AGGR_SUM;
646 } else if (!strcasecmp(c->argv[j]->ptr,"min")) {
647 aggregate = REDIS_AGGR_MIN;
648 } else if (!strcasecmp(c->argv[j]->ptr,"max")) {
649 aggregate = REDIS_AGGR_MAX;
650 } else {
651 zfree(src);
652 addReply(c,shared.syntaxerr);
653 return;
654 }
655 j++; remaining--;
656 } else {
657 zfree(src);
658 addReply(c,shared.syntaxerr);
659 return;
660 }
661 }
662 }
663
664 /* sort sets from the smallest to largest, this will improve our
665 * algorithm's performance */
666 qsort(src,setnum,sizeof(zsetopsrc),qsortCompareZsetopsrcByCardinality);
667
668 dstobj = createZsetObject();
669 dstzset = dstobj->ptr;
670
671 if (op == REDIS_OP_INTER) {
672 /* skip going over all entries if the smallest zset is NULL or empty */
673 if (src[0].dict && dictSize(src[0].dict) > 0) {
674 /* precondition: as src[0].dict is non-empty and the zsets are ordered
675 * from small to large, all src[i > 0].dict are non-empty too */
676 di = dictGetIterator(src[0].dict);
677 while((de = dictNext(di)) != NULL) {
678 double score, value;
679
680 score = src[0].weight * zunionInterDictValue(de);
681 for (j = 1; j < setnum; j++) {
682 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
683 if (other) {
684 value = src[j].weight * zunionInterDictValue(other);
685 zunionInterAggregate(&score,value,aggregate);
686 } else {
687 break;
688 }
689 }
690
691 /* Only continue when present in every source dict. */
692 if (j == setnum) {
693 robj *o = dictGetEntryKey(de);
694 znode = zslInsert(dstzset->zsl,score,o);
695 incrRefCount(o); /* added to skiplist */
696 dictAdd(dstzset->dict,o,&znode->score);
697 incrRefCount(o); /* added to dictionary */
698 }
699 }
700 dictReleaseIterator(di);
701 }
702 } else if (op == REDIS_OP_UNION) {
703 for (i = 0; i < setnum; i++) {
704 if (!src[i].dict) continue;
705
706 di = dictGetIterator(src[i].dict);
707 while((de = dictNext(di)) != NULL) {
708 double score, value;
709
710 /* skip key when already processed */
711 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL)
712 continue;
713
714 /* initialize score */
715 score = src[i].weight * zunionInterDictValue(de);
716
717 /* because the zsets are sorted by size, its only possible
718 * for sets at larger indices to hold this entry */
719 for (j = (i+1); j < setnum; j++) {
720 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
721 if (other) {
722 value = src[j].weight * zunionInterDictValue(other);
723 zunionInterAggregate(&score,value,aggregate);
724 }
725 }
726
727 robj *o = dictGetEntryKey(de);
728 znode = zslInsert(dstzset->zsl,score,o);
729 incrRefCount(o); /* added to skiplist */
730 dictAdd(dstzset->dict,o,&znode->score);
731 incrRefCount(o); /* added to dictionary */
732 }
733 dictReleaseIterator(di);
734 }
735 } else {
736 /* unknown operator */
737 redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION);
738 }
739
740 if (dbDelete(c->db,dstkey)) {
741 touchWatchedKey(c->db,dstkey);
742 touched = 1;
743 server.dirty++;
744 }
745 if (dstzset->zsl->length) {
746 dbAdd(c->db,dstkey,dstobj);
747 addReplyLongLong(c, dstzset->zsl->length);
748 if (!touched) touchWatchedKey(c->db,dstkey);
749 server.dirty++;
750 } else {
751 decrRefCount(dstobj);
752 addReply(c, shared.czero);
753 }
754 zfree(src);
755 }
756
757 void zunionstoreCommand(redisClient *c) {
758 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION);
759 }
760
761 void zinterstoreCommand(redisClient *c) {
762 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
763 }
764
765 void zrangeGenericCommand(redisClient *c, int reverse) {
766 robj *o;
767 long start;
768 long end;
769 int withscores = 0;
770 int llen;
771 int rangelen, j;
772 zset *zsetobj;
773 zskiplist *zsl;
774 zskiplistNode *ln;
775 robj *ele;
776
777 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
778 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
779
780 if (c->argc == 5 && !strcasecmp(c->argv[4]->ptr,"withscores")) {
781 withscores = 1;
782 } else if (c->argc >= 5) {
783 addReply(c,shared.syntaxerr);
784 return;
785 }
786
787 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) == NULL
788 || checkType(c,o,REDIS_ZSET)) return;
789 zsetobj = o->ptr;
790 zsl = zsetobj->zsl;
791 llen = zsl->length;
792
793 /* convert negative indexes */
794 if (start < 0) start = llen+start;
795 if (end < 0) end = llen+end;
796 if (start < 0) start = 0;
797
798 /* Invariant: start >= 0, so this test will be true when end < 0.
799 * The range is empty when start > end or start >= length. */
800 if (start > end || start >= llen) {
801 addReply(c,shared.emptymultibulk);
802 return;
803 }
804 if (end >= llen) end = llen-1;
805 rangelen = (end-start)+1;
806
807 /* check if starting point is trivial, before searching
808 * the element in log(N) time */
809 if (reverse) {
810 ln = start == 0 ? zsl->tail : zslistTypeGetElementByRank(zsl, llen-start);
811 } else {
812 ln = start == 0 ?
813 zsl->header->level[0].forward : zslistTypeGetElementByRank(zsl, start+1);
814 }
815
816 /* Return the result in form of a multi-bulk reply */
817 addReplyMultiBulkLen(c,withscores ? (rangelen*2) : rangelen);
818 for (j = 0; j < rangelen; j++) {
819 ele = ln->obj;
820 addReplyBulk(c,ele);
821 if (withscores)
822 addReplyDouble(c,ln->score);
823 ln = reverse ? ln->backward : ln->level[0].forward;
824 }
825 }
826
827 void zrangeCommand(redisClient *c) {
828 zrangeGenericCommand(c,0);
829 }
830
831 void zrevrangeCommand(redisClient *c) {
832 zrangeGenericCommand(c,1);
833 }
834
835 /* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE and ZCOUNT.
836 * If "justcount", only the number of elements in the range is returned. */
837 void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) {
838 zrangespec range;
839 robj *o, *emptyreply;
840 zset *zsetobj;
841 zskiplist *zsl;
842 zskiplistNode *ln;
843 int offset = 0, limit = -1;
844 int withscores = 0;
845 unsigned long rangelen = 0;
846 void *replylen = NULL;
847
848 /* Parse the range arguments. */
849 if (zslParseRange(c->argv[2],c->argv[3],&range) != REDIS_OK) {
850 addReplyError(c,"min or max is not a double");
851 return;
852 }
853
854 /* Parse optional extra arguments. Note that ZCOUNT will exactly have
855 * 4 arguments, so we'll never enter the following code path. */
856 if (c->argc > 4) {
857 int remaining = c->argc - 4;
858 int pos = 4;
859
860 while (remaining) {
861 if (remaining >= 1 && !strcasecmp(c->argv[pos]->ptr,"withscores")) {
862 pos++; remaining--;
863 withscores = 1;
864 } else if (remaining >= 3 && !strcasecmp(c->argv[pos]->ptr,"limit")) {
865 offset = atoi(c->argv[pos+1]->ptr);
866 limit = atoi(c->argv[pos+2]->ptr);
867 pos += 3; remaining -= 3;
868 } else {
869 addReply(c,shared.syntaxerr);
870 return;
871 }
872 }
873 }
874
875 /* Ok, lookup the key and get the range */
876 emptyreply = justcount ? shared.czero : shared.emptymultibulk;
877 if ((o = lookupKeyReadOrReply(c,c->argv[1],emptyreply)) == NULL ||
878 checkType(c,o,REDIS_ZSET)) return;
879 zsetobj = o->ptr;
880 zsl = zsetobj->zsl;
881
882 /* If reversed, assume the elements are sorted from high to low score. */
883 ln = zslFirstWithScore(zsl,range.min);
884 if (reverse) {
885 /* If range.min is out of range, ln will be NULL and we need to use
886 * the tail of the skiplist as first node of the range. */
887 if (ln == NULL) ln = zsl->tail;
888
889 /* zslFirstWithScore returns the first element with where with
890 * score >= range.min, so backtrack to make sure the element we use
891 * here has score <= range.min. */
892 while (ln && ln->score > range.min) ln = ln->backward;
893
894 /* Move to the right element according to the range spec. */
895 if (range.minex) {
896 /* Find last element with score < range.min */
897 while (ln && ln->score == range.min) ln = ln->backward;
898 } else {
899 /* Find last element with score <= range.min */
900 while (ln && ln->level[0].forward &&
901 ln->level[0].forward->score == range.min)
902 ln = ln->level[0].forward;
903 }
904 } else {
905 if (range.minex) {
906 /* Find first element with score > range.min */
907 while (ln && ln->score == range.min) ln = ln->level[0].forward;
908 }
909 }
910
911 /* No "first" element in the specified interval. */
912 if (ln == NULL) {
913 addReply(c,emptyreply);
914 return;
915 }
916
917 /* We don't know in advance how many matching elements there
918 * are in the list, so we push this object that will represent
919 * the multi-bulk length in the output buffer, and will "fix"
920 * it later */
921 if (!justcount)
922 replylen = addDeferredMultiBulkLength(c);
923
924 /* If there is an offset, just traverse the number of elements without
925 * checking the score because that is done in the next loop. */
926 while(ln && offset--) {
927 if (reverse)
928 ln = ln->backward;
929 else
930 ln = ln->level[0].forward;
931 }
932
933 while (ln && limit--) {
934 /* Check if this this element is in range. */
935 if (reverse) {
936 if (range.maxex) {
937 /* Element should have score > range.max */
938 if (ln->score <= range.max) break;
939 } else {
940 /* Element should have score >= range.max */
941 if (ln->score < range.max) break;
942 }
943 } else {
944 if (range.maxex) {
945 /* Element should have score < range.max */
946 if (ln->score >= range.max) break;
947 } else {
948 /* Element should have score <= range.max */
949 if (ln->score > range.max) break;
950 }
951 }
952
953 /* Do our magic */
954 rangelen++;
955 if (!justcount) {
956 addReplyBulk(c,ln->obj);
957 if (withscores)
958 addReplyDouble(c,ln->score);
959 }
960
961 if (reverse)
962 ln = ln->backward;
963 else
964 ln = ln->level[0].forward;
965 }
966
967 if (justcount) {
968 addReplyLongLong(c,(long)rangelen);
969 } else {
970 setDeferredMultiBulkLength(c,replylen,
971 withscores ? (rangelen*2) : rangelen);
972 }
973 }
974
975 void zrangebyscoreCommand(redisClient *c) {
976 genericZrangebyscoreCommand(c,0,0);
977 }
978
979 void zrevrangebyscoreCommand(redisClient *c) {
980 genericZrangebyscoreCommand(c,1,0);
981 }
982
983 void zcountCommand(redisClient *c) {
984 genericZrangebyscoreCommand(c,0,1);
985 }
986
987 void zcardCommand(redisClient *c) {
988 robj *o;
989 zset *zs;
990
991 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
992 checkType(c,o,REDIS_ZSET)) return;
993
994 zs = o->ptr;
995 addReplyLongLong(c,zs->zsl->length);
996 }
997
998 void zscoreCommand(redisClient *c) {
999 robj *o;
1000 zset *zs;
1001 dictEntry *de;
1002
1003 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1004 checkType(c,o,REDIS_ZSET)) return;
1005
1006 zs = o->ptr;
1007 de = dictFind(zs->dict,c->argv[2]);
1008 if (!de) {
1009 addReply(c,shared.nullbulk);
1010 } else {
1011 double *score = dictGetEntryVal(de);
1012
1013 addReplyDouble(c,*score);
1014 }
1015 }
1016
1017 void zrankGenericCommand(redisClient *c, int reverse) {
1018 robj *o;
1019 zset *zs;
1020 zskiplist *zsl;
1021 dictEntry *de;
1022 unsigned long rank;
1023 double *score;
1024
1025 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1026 checkType(c,o,REDIS_ZSET)) return;
1027
1028 zs = o->ptr;
1029 zsl = zs->zsl;
1030 de = dictFind(zs->dict,c->argv[2]);
1031 if (!de) {
1032 addReply(c,shared.nullbulk);
1033 return;
1034 }
1035
1036 score = dictGetEntryVal(de);
1037 rank = zslistTypeGetRank(zsl, *score, c->argv[2]);
1038 if (rank) {
1039 if (reverse) {
1040 addReplyLongLong(c, zsl->length - rank);
1041 } else {
1042 addReplyLongLong(c, rank-1);
1043 }
1044 } else {
1045 addReply(c,shared.nullbulk);
1046 }
1047 }
1048
1049 void zrankCommand(redisClient *c) {
1050 zrankGenericCommand(c, 0);
1051 }
1052
1053 void zrevrankCommand(redisClient *c) {
1054 zrankGenericCommand(c, 1);
1055 }