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