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