Add GT and LT options to ZADD for conditional score updates (#7818)

Co-authored-by: Alex Ronke <w.alex.ronke@gmail.com>
This commit is contained in:
alexronke-channeladvisor 2020-09-23 14:56:16 -04:00 committed by GitHub
parent a735bf5c2a
commit 66a13ccbdf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 98 additions and 2 deletions

1
.gitignore vendored
View File

@ -37,3 +37,4 @@ Makefile.dep
.ccls
.ccls-cache/*
compile_commands.json
redis.code-workspace

View File

@ -2380,6 +2380,8 @@ int RM_ZsetAddFlagsToCoreFlags(int flags) {
int retflags = 0;
if (flags & REDISMODULE_ZADD_XX) retflags |= ZADD_XX;
if (flags & REDISMODULE_ZADD_NX) retflags |= ZADD_NX;
if (flags & REDISMODULE_ZADD_GT) retflags |= ZADD_GT;
if (flags & REDISMODULE_ZADD_LT) retflags |= ZADD_LT;
return retflags;
}
@ -2406,6 +2408,10 @@ int RM_ZsetAddFlagsFromCoreFlags(int flags) {
*
* REDISMODULE_ZADD_XX: Element must already exist. Do nothing otherwise.
* REDISMODULE_ZADD_NX: Element must not exist. Do nothing otherwise.
* REDISMODULE_ZADD_GT: If element exists, new score must be greater than the current score.
* Do nothing otherwise. Can optionally be combined with XX.
* REDISMODULE_ZADD_LT: If element exists, new score must be less than the current score.
* Do nothing otherwise. Can optionally be combined with XX.
*
* The output flags are:
*

View File

@ -55,6 +55,8 @@
#define REDISMODULE_ZADD_ADDED (1<<2)
#define REDISMODULE_ZADD_UPDATED (1<<3)
#define REDISMODULE_ZADD_NOP (1<<4)
#define REDISMODULE_ZADD_GT (1<<5)
#define REDISMODULE_ZADD_LT (1<<6)
/* Hash API flags. */
#define REDISMODULE_HASH_NONE 0

View File

@ -1945,6 +1945,8 @@ void addACLLogEntry(client *c, int reason, int keypos, sds username);
#define ZADD_INCR (1<<0) /* Increment the score instead of setting it. */
#define ZADD_NX (1<<1) /* Don't touch elements not already existing. */
#define ZADD_XX (1<<2) /* Only touch elements already existing. */
#define ZADD_GT (1<<7) /* Only update existing when new scores are higher. */
#define ZADD_LT (1<<8) /* Only update existing when new scores are lower. */
/* Output flags. */
#define ZADD_NOP (1<<3) /* Operation not performed because of conditionals.*/

View File

@ -1283,6 +1283,10 @@ int zsetScore(robj *zobj, sds member, double *score) {
* assume 0 as previous score.
* ZADD_NX: Perform the operation only if the element does not exist.
* ZADD_XX: Perform the operation only if the element already exist.
* ZADD_GT: Perform the operation on existing elements only if the new score is
* greater than the current score.
* ZADD_LT: Perform the operation on existing elements only if the new score is
* less than the current score.
*
* When ZADD_INCR is used, the new score of the element is stored in
* '*newscore' if 'newscore' is not NULL.
@ -1317,6 +1321,8 @@ int zsetAdd(robj *zobj, double score, sds ele, int *flags, double *newscore) {
int incr = (*flags & ZADD_INCR) != 0;
int nx = (*flags & ZADD_NX) != 0;
int xx = (*flags & ZADD_XX) != 0;
int gt = (*flags & ZADD_GT) != 0;
int lt = (*flags & ZADD_LT) != 0;
*flags = 0; /* We'll return our response flags. */
double curscore;
@ -1348,7 +1354,12 @@ int zsetAdd(robj *zobj, double score, sds ele, int *flags, double *newscore) {
}
/* Remove and re-insert when score changed. */
if (score != curscore) {
if (score != curscore &&
/* LT? Only update if score is less than current. */
(!lt || score < curscore) &&
/* GT? Only update if score is greater than current. */
(!gt || score > curscore))
{
zobj->ptr = zzlDelete(zobj->ptr,eptr);
zobj->ptr = zzlInsert(zobj->ptr,ele,score);
*flags |= ZADD_UPDATED;
@ -1393,7 +1404,12 @@ int zsetAdd(robj *zobj, double score, sds ele, int *flags, double *newscore) {
}
/* Remove and re-insert when score changes. */
if (score != curscore) {
if (score != curscore &&
/* LT? Only update if score is less than current. */
(!lt || score < curscore) &&
/* GT? Only update if score is greater than current. */
(!gt || score > curscore))
{
znode = zslUpdateScore(zs->zsl,curscore,ele,score);
/* Note that we did not removed the original element from
* the hash table representing the sorted set, so we just
@ -1555,6 +1571,8 @@ void zaddGenericCommand(client *c, int flags) {
else if (!strcasecmp(opt,"xx")) flags |= ZADD_XX;
else if (!strcasecmp(opt,"ch")) flags |= ZADD_CH;
else if (!strcasecmp(opt,"incr")) flags |= ZADD_INCR;
else if (!strcasecmp(opt,"gt")) flags |= ZADD_GT;
else if (!strcasecmp(opt,"lt")) flags |= ZADD_LT;
else break;
scoreidx++;
}
@ -1564,6 +1582,8 @@ void zaddGenericCommand(client *c, int flags) {
int nx = (flags & ZADD_NX) != 0;
int xx = (flags & ZADD_XX) != 0;
int ch = (flags & ZADD_CH) != 0;
int gt = (flags & ZADD_GT) != 0;
int lt = (flags & ZADD_LT) != 0;
/* After the options, we expect to have an even number of args, since
* we expect any number of score-element pairs. */
@ -1580,6 +1600,13 @@ void zaddGenericCommand(client *c, int flags) {
"XX and NX options at the same time are not compatible");
return;
}
if ((gt && nx) || (lt && nx) || (gt && lt)) {
addReplyError(c,
"GT, LT, and/or NX options at the same time are not compatible");
return;
}
/* Note that XX is compatible with either GT or LT */
if (incr && elements > 1) {
addReplyError(c,

View File

@ -78,6 +78,46 @@ start_server {tags {"zset"}} {
assert {[r zscore ztmp y] == 21}
}
test "ZADD GT updates existing elements when new scores are greater" {
r del ztmp
r zadd ztmp 10 x 20 y 30 z
assert {[r zadd ztmp gt ch 5 foo 11 x 21 y 29 z] == 3}
assert {[r zcard ztmp] == 4}
assert {[r zscore ztmp x] == 11}
assert {[r zscore ztmp y] == 21}
assert {[r zscore ztmp z] == 30}
}
test "ZADD LT updates existing elements when new scores are lower" {
r del ztmp
r zadd ztmp 10 x 20 y 30 z
assert {[r zadd ztmp lt ch 5 foo 11 x 21 y 29 z] == 2}
assert {[r zcard ztmp] == 4}
assert {[r zscore ztmp x] == 10}
assert {[r zscore ztmp y] == 20}
assert {[r zscore ztmp z] == 29}
}
test "ZADD GT XX updates existing elements when new scores are greater and skips new elements" {
r del ztmp
r zadd ztmp 10 x 20 y 30 z
assert {[r zadd ztmp gt xx ch 5 foo 11 x 21 y 29 z] == 2}
assert {[r zcard ztmp] == 3}
assert {[r zscore ztmp x] == 11}
assert {[r zscore ztmp y] == 21}
assert {[r zscore ztmp z] == 30}
}
test "ZADD LT XX updates existing elements when new scores are lower and skips new elements" {
r del ztmp
r zadd ztmp 10 x 20 y 30 z
assert {[r zadd ztmp lt xx ch 5 foo 11 x 21 y 29 z] == 1}
assert {[r zcard ztmp] == 3}
assert {[r zscore ztmp x] == 10}
assert {[r zscore ztmp y] == 20}
assert {[r zscore ztmp z] == 29}
}
test "ZADD XX and NX are not compatible" {
r del ztmp
catch {r zadd ztmp xx nx 10 x} err
@ -100,6 +140,24 @@ start_server {tags {"zset"}} {
assert {[r zscore ztmp b] == 200}
}
test "ZADD GT and NX are not compatible" {
r del ztmp
catch {r zadd ztmp gt nx 10 x} err
set err
} {ERR*}
test "ZADD LT and NX are not compatible" {
r del ztmp
catch {r zadd ztmp lt nx 10 x} err
set err
} {ERR*}
test "ZADD LT and GT are not compatible" {
r del ztmp
catch {r zadd ztmp lt gt 10 x} err
set err
} {ERR*}
test "ZADD INCR works like ZINCRBY" {
r del ztmp
r zadd ztmp 10 x 20 y 30 z