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