Skip to content

Add ZRANK, ZREVRANK WITHSCORE #2531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Apr 19, 2023
Merged
60 changes: 60 additions & 0 deletions command.go
Original file line number Diff line number Diff line change
Expand Up @@ -4724,3 +4724,63 @@ func (cmd *ClusterShardsCmd) readReply(rd *proto.Reader) error {

return nil
}

// -----------------------------------------

type RankScore struct {
Rank int64
Score float64
}

type RankWithScoreCmd struct {
baseCmd

val RankScore
}

var _ Cmder = (*RankWithScoreCmd)(nil)

func NewRankWithScoreCmd(ctx context.Context, args ...interface{}) *RankWithScoreCmd {
return &RankWithScoreCmd{
baseCmd: baseCmd{
ctx: ctx,
args: args,
},
}
}

func (cmd *RankWithScoreCmd) SetVal(val RankScore) {
cmd.val = val
}

func (cmd *RankWithScoreCmd) Val() RankScore {
return cmd.val
}

func (cmd *RankWithScoreCmd) Result() (RankScore, error) {
return cmd.val, cmd.err
}

func (cmd *RankWithScoreCmd) String() string {
return cmdString(cmd, cmd.val)
}

func (cmd *RankWithScoreCmd) readReply(rd *proto.Reader) error {
if err := rd.ReadFixedArrayLen(2); err != nil {
return err
}

rank, err := rd.ReadInt()
if err != nil {
return err
}

score, err := rd.ReadFloat()
if err != nil {
return err
}

cmd.val = RankScore{Rank: rank, Score: score}

return nil
}
18 changes: 18 additions & 0 deletions commands.go
Original file line number Diff line number Diff line change
Expand Up @@ -373,6 +373,7 @@ type Cmdable interface {
ZRangeArgsWithScores(ctx context.Context, z ZRangeArgs) *ZSliceCmd
ZRangeStore(ctx context.Context, dst string, z ZRangeArgs) *IntCmd
ZRank(ctx context.Context, key, member string) *IntCmd
ZRankWithScore(ctx context.Context, key, member string) *RankWithScoreCmd
ZRem(ctx context.Context, key string, members ...interface{}) *IntCmd
ZRemRangeByRank(ctx context.Context, key string, start, stop int64) *IntCmd
ZRemRangeByScore(ctx context.Context, key, min, max string) *IntCmd
Expand All @@ -383,6 +384,7 @@ type Cmdable interface {
ZRevRangeByLex(ctx context.Context, key string, opt *ZRangeBy) *StringSliceCmd
ZRevRangeByScoreWithScores(ctx context.Context, key string, opt *ZRangeBy) *ZSliceCmd
ZRevRank(ctx context.Context, key, member string) *IntCmd
ZRevRankWithScore(ctx context.Context, key, member string) *RankWithScoreCmd
ZScore(ctx context.Context, key, member string) *FloatCmd
ZUnionStore(ctx context.Context, dest string, store *ZStore) *IntCmd
ZRandMember(ctx context.Context, key string, count int) *StringSliceCmd
Expand Down Expand Up @@ -2884,6 +2886,14 @@ func (c cmdable) ZRank(ctx context.Context, key, member string) *IntCmd {
return cmd
}

// ZRankWithScore according to the Redis documentation, if member does not exist
// in the sorted set or key does not exist, it will return a redis.Nil error.
func (c cmdable) ZRankWithScore(ctx context.Context, key, member string) *RankWithScoreCmd {
cmd := NewRankWithScoreCmd(ctx, "zrank", key, member, "withscore")
_ = c(ctx, cmd)
return cmd
}

func (c cmdable) ZRem(ctx context.Context, key string, members ...interface{}) *IntCmd {
args := make([]interface{}, 2, 2+len(members))
args[0] = "zrem"
Expand Down Expand Up @@ -2924,6 +2934,8 @@ func (c cmdable) ZRevRange(ctx context.Context, key string, start, stop int64) *
return cmd
}

// ZRevRangeWithScores according to the Redis documentation, if member does not exist
// in the sorted set or key does not exist, it will return a redis.Nil error.
func (c cmdable) ZRevRangeWithScores(ctx context.Context, key string, start, stop int64) *ZSliceCmd {
cmd := NewZSliceCmd(ctx, "zrevrange", key, start, stop, "withscores")
_ = c(ctx, cmd)
Expand Down Expand Up @@ -2974,6 +2986,12 @@ func (c cmdable) ZRevRank(ctx context.Context, key, member string) *IntCmd {
return cmd
}

func (c cmdable) ZRevRankWithScore(ctx context.Context, key, member string) *RankWithScoreCmd {
cmd := NewRankWithScoreCmd(ctx, "zrevrank", key, member, "withscore")
_ = c(ctx, cmd)
return cmd
}

func (c cmdable) ZScore(ctx context.Context, key, member string) *FloatCmd {
cmd := NewFloatCmd(ctx, "zscore", key, member)
_ = c(ctx, cmd)
Expand Down
50 changes: 50 additions & 0 deletions commands_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4737,6 +4737,31 @@ var _ = Describe("Commands", func() {
Expect(zRank.Val()).To(Equal(int64(0)))
})

It("should ZRankWithScore", func() {
err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 2, Member: "two"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 3, Member: "three"}).Err()
Expect(err).NotTo(HaveOccurred())

zRankWithScore := client.ZRankWithScore(ctx, "zset", "one")
Expect(zRankWithScore.Err()).NotTo(HaveOccurred())
Expect(zRankWithScore.Result()).To(Equal(redis.RankScore{Rank: 0, Score: 1}))

zRankWithScore = client.ZRankWithScore(ctx, "zset", "two")
Expect(zRankWithScore.Err()).NotTo(HaveOccurred())
Expect(zRankWithScore.Result()).To(Equal(redis.RankScore{Rank: 1, Score: 2}))

zRankWithScore = client.ZRankWithScore(ctx, "zset", "three")
Expect(zRankWithScore.Err()).NotTo(HaveOccurred())
Expect(zRankWithScore.Result()).To(Equal(redis.RankScore{Rank: 2, Score: 3}))

zRankWithScore = client.ZRankWithScore(ctx, "zset", "four")
Expect(zRankWithScore.Err()).To(HaveOccurred())
Expect(zRankWithScore.Err()).To(Equal(redis.Nil))
})

It("should ZRem", func() {
err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
Expect(err).NotTo(HaveOccurred())
Expand Down Expand Up @@ -5008,6 +5033,31 @@ var _ = Describe("Commands", func() {
Expect(zRevRank.Val()).To(Equal(int64(0)))
})

It("should ZRevRankWithScore", func() {
err := client.ZAdd(ctx, "zset", redis.Z{Score: 1, Member: "one"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 2, Member: "two"}).Err()
Expect(err).NotTo(HaveOccurred())
err = client.ZAdd(ctx, "zset", redis.Z{Score: 3, Member: "three"}).Err()
Expect(err).NotTo(HaveOccurred())

zRevRankWithScore := client.ZRevRankWithScore(ctx, "zset", "one")
Expect(zRevRankWithScore.Err()).NotTo(HaveOccurred())
Expect(zRevRankWithScore.Result()).To(Equal(redis.RankScore{Rank: 2, Score: 1}))

zRevRankWithScore = client.ZRevRankWithScore(ctx, "zset", "two")
Expect(zRevRankWithScore.Err()).NotTo(HaveOccurred())
Expect(zRevRankWithScore.Result()).To(Equal(redis.RankScore{Rank: 1, Score: 2}))

zRevRankWithScore = client.ZRevRankWithScore(ctx, "zset", "three")
Expect(zRevRankWithScore.Err()).NotTo(HaveOccurred())
Expect(zRevRankWithScore.Result()).To(Equal(redis.RankScore{Rank: 0, Score: 3}))

zRevRankWithScore = client.ZRevRankWithScore(ctx, "zset", "four")
Expect(zRevRankWithScore.Err()).To(HaveOccurred())
Expect(zRevRankWithScore.Err()).To(Equal(redis.Nil))
})

It("should ZScore", func() {
zAdd := client.ZAdd(ctx, "zset", redis.Z{Score: 1.001, Member: "one"})
Expect(zAdd.Err()).NotTo(HaveOccurred())
Expand Down