| 4015 | #define ZRANDMEMBER_RANDOM_SAMPLE_LIMIT 1000 |
| 4016 | |
| 4017 | void zrandmemberWithCountCommand(client *c, long l, int withscores) { |
| 4018 | unsigned long count, size; |
| 4019 | int uniq = 1; |
| 4020 | robj *zsetobj; |
| 4021 | |
| 4022 | if ((zsetobj = lookupKeyReadOrReply(c, c->argv[1], shared.emptyarray)) |
| 4023 | == NULL || checkType(c, zsetobj, OBJ_ZSET)) return; |
| 4024 | size = zsetLength(zsetobj); |
| 4025 | |
| 4026 | if(l >= 0) { |
| 4027 | count = (unsigned long) l; |
| 4028 | } else { |
| 4029 | count = -l; |
| 4030 | uniq = 0; |
| 4031 | } |
| 4032 | |
| 4033 | /* If count is zero, serve it ASAP to avoid special cases later. */ |
| 4034 | if (count == 0) { |
| 4035 | addReply(c,shared.emptyarray); |
| 4036 | return; |
| 4037 | } |
| 4038 | |
| 4039 | /* CASE 1: The count was negative, so the extraction method is just: |
| 4040 | * "return N random elements" sampling the whole set every time. |
| 4041 | * This case is trivial and can be served without auxiliary data |
| 4042 | * structures. This case is the only one that also needs to return the |
| 4043 | * elements in random order. */ |
| 4044 | if (!uniq || count == 1) { |
| 4045 | if (withscores && c->resp == 2) |
| 4046 | addReplyArrayLen(c, count*2); |
| 4047 | else |
| 4048 | addReplyArrayLen(c, count); |
| 4049 | if (zsetobj->encoding == OBJ_ENCODING_SKIPLIST) { |
| 4050 | zset *zs = zsetobj->ptr; |
| 4051 | while (count--) { |
| 4052 | dictEntry *de = dictGetFairRandomKey(zs->dict); |
| 4053 | sds key = dictGetKey(de); |
| 4054 | if (withscores && c->resp > 2) |
| 4055 | addReplyArrayLen(c,2); |
| 4056 | addReplyBulkCBuffer(c, key, sdslen(key)); |
| 4057 | if (withscores) |
| 4058 | addReplyDouble(c, *(double*)dictGetVal(de)); |
| 4059 | } |
| 4060 | } else if (zsetobj->encoding == OBJ_ENCODING_ZIPLIST) { |
| 4061 | ziplistEntry *keys, *vals = NULL; |
| 4062 | unsigned long limit, sample_count; |
| 4063 | limit = count > ZRANDMEMBER_RANDOM_SAMPLE_LIMIT ? ZRANDMEMBER_RANDOM_SAMPLE_LIMIT : count; |
| 4064 | keys = zmalloc(sizeof(ziplistEntry)*limit); |
| 4065 | if (withscores) |
| 4066 | vals = zmalloc(sizeof(ziplistEntry)*limit); |
| 4067 | while (count) { |
| 4068 | sample_count = count > limit ? limit : count; |
| 4069 | count -= sample_count; |
| 4070 | ziplistRandomPairs(zsetobj->ptr, sample_count, keys, vals); |
| 4071 | zarndmemberReplyWithZiplist(c, sample_count, keys, vals); |
| 4072 | } |
| 4073 | zfree(keys); |
| 4074 | zfree(vals); |
no test coverage detected