Merge branch 'zset_score'

This commit is contained in:
siddontang 2014-06-05 14:21:50 +08:00
parent bd38e693f2
commit 5ea473b4fe
4 changed files with 76 additions and 28 deletions

View File

@ -37,4 +37,5 @@ var (
ErrKeySize = errors.New("invalid key size")
ErrHashFieldSize = errors.New("invalid hash field size")
ErrZSetMemberSize = errors.New("invalid zset member size")
ErrScoreMiss = errors.New("zset score miss")
)

View File

@ -10,6 +10,7 @@ import (
const (
MinScore int64 = -1<<63 + 1
MaxScore int64 = 1<<63 - 1
InvalidScore int64 = -1 << 63
)
type ScorePair struct {
@ -311,19 +312,25 @@ func (db *DB) ZCard(key []byte) (int64, error) {
return Int64(db.db.Get(sk))
}
func (db *DB) ZScore(key []byte, member []byte) ([]byte, error) {
func (db *DB) ZScore(key []byte, member []byte) (int64, error) {
if err := checkZSetKMSize(key, member); err != nil {
return nil, err
return InvalidScore, err
}
var score int64 = InvalidScore
k := db.zEncodeSetKey(key, member)
score, err := Int64(db.db.Get(k))
if err != nil {
return nil, err
if v, err := db.db.Get(k); err != nil {
return InvalidScore, err
} else if v == nil {
return InvalidScore, ErrScoreMiss
} else {
if score, err = Int64(v, nil); err != nil {
return InvalidScore, err
}
}
return StrPutInt64(score), nil
return score, nil
}
func (db *DB) ZRem(key []byte, members ...[]byte) (int64, error) {
@ -356,9 +363,9 @@ func (db *DB) ZRem(key []byte, members ...[]byte) (int64, error) {
return num, err
}
func (db *DB) ZIncrBy(key []byte, delta int64, member []byte) ([]byte, error) {
func (db *DB) ZIncrBy(key []byte, delta int64, member []byte) (int64, error) {
if err := checkZSetKMSize(key, member); err != nil {
return nil, err
return InvalidScore, err
}
t := db.zsetTx
@ -371,18 +378,17 @@ func (db *DB) ZIncrBy(key []byte, delta int64, member []byte) ([]byte, error) {
v, err := db.db.Get(ek)
if err != nil {
return nil, err
return InvalidScore, err
} else if v != nil {
if s, err := Int64(v, err); err != nil {
return nil, err
return InvalidScore, err
} else {
sk := db.zEncodeScoreKey(key, member, s)
t.Delete(sk)
score = s + delta
if score >= MaxScore || score <= MinScore {
return nil, errScoreOverflow
return InvalidScore, errScoreOverflow
}
}
} else {
@ -395,7 +401,7 @@ func (db *DB) ZIncrBy(key []byte, delta int64, member []byte) ([]byte, error) {
t.Put(sk, []byte{})
err = t.Commit()
return StrPutInt64(score), err
return score, err
}
func (db *DB) ZCount(key []byte, min int64, max int64) (int64, error) {
@ -563,10 +569,10 @@ func (db *DB) zRange(key []byte, min int64, max int64, withScores bool, offset i
continue
}
v = append(v, m)
if withScores {
v = append(v, StrPutInt64(s))
v = append(v, m, s)
} else {
v = append(v, m)
}
}

View File

@ -71,6 +71,16 @@ func TestDBZSet(t *testing.T) {
t.Fatal(n)
}
if s, err := db.ZScore(key, bin("d")); err != nil {
t.Fatal(err)
} else if s != 3 {
t.Fatal(s)
}
if s, err := db.ZScore(key, bin("zzz")); err != ErrScoreMiss || s != InvalidScore {
t.Fatal(fmt.Sprintf("s=[%d] err=[%s]", s, err))
}
// {c':2, 'd':3}
if n, err := db.ZRem(key, bin("a"), bin("b")); err != nil {
t.Fatal(err)
@ -179,8 +189,10 @@ func TestZSetOrder(t *testing.T) {
}
// {'a':0, 'b':1, 'c':2, 'd':999, 'e':6, 'f':5}
if _, err := db.ZIncrBy(key, 2, bin("e")); err != nil {
if s, err := db.ZIncrBy(key, 2, bin("e")); err != nil {
t.Fatal(err)
} else if s != 6 {
t.Fatal(s)
}
if pos, _ := db.ZRank(key, bin("e")); int(pos) != 4 {
@ -190,6 +202,19 @@ func TestZSetOrder(t *testing.T) {
if pos, _ := db.ZRevRank(key, bin("e")); int(pos) != 1 {
t.Fatal(pos)
}
if datas, _ := db.ZRange(key, 0, endPos, true); len(datas) != 12 {
t.Fatal(len(datas))
} else {
scores := []int64{0, 1, 2, 5, 6, 999}
for i := 1; i < len(datas); i += 2 {
if s, ok := datas[i].(int64); !ok || s != scores[(i-1)/2] {
t.Fatal(fmt.Sprintf("[%d]=%d", i, datas[i]))
}
}
}
return
}
func TestDBZScan(t *testing.T) {

View File

@ -66,10 +66,14 @@ func zscoreCommand(c *client) error {
return ErrCmdParams
}
if v, err := c.db.ZScore(args[0], args[1]); err != nil {
return err
if s, err := c.db.ZScore(args[0], args[1]); err != nil {
if err == ledis.ErrScoreMiss {
c.writeBulk(nil)
} else {
c.writeBulk(v)
return err
}
} else {
c.writeBulk(ledis.StrPutInt64(s))
}
return nil
@ -106,7 +110,7 @@ func zincrbyCommand(c *client) error {
if v, err := c.db.ZIncrBy(key, delta, args[2]); err != nil {
return err
} else {
c.writeBulk(v)
c.writeBulk(ledis.StrPutInt64(v))
}
return nil
@ -308,10 +312,16 @@ func zrangeGeneric(c *client, reverse bool) error {
withScores = true
}
if v, err := c.db.ZRangeGeneric(key, start, stop, withScores, reverse); err != nil {
if datas, err := c.db.ZRangeGeneric(key, start, stop, withScores, reverse); err != nil {
return err
} else {
c.writeArray(v)
if withScores {
for i := len(datas) - 1; i > 0; i -= 2 {
v, _ := datas[i].(int64)
datas[i] = ledis.StrPutInt64(v)
}
}
c.writeArray(datas)
}
return nil
}
@ -373,10 +383,16 @@ func zrangebyscoreGeneric(c *client, reverse bool) error {
return nil
}
if v, err := c.db.ZRangeByScoreGeneric(key, min, max, withScores, offset, count, reverse); err != nil {
if datas, err := c.db.ZRangeByScoreGeneric(key, min, max, withScores, offset, count, reverse); err != nil {
return err
} else {
c.writeArray(v)
if withScores {
for i := len(datas) - 1; i > 0; i -= 2 {
v, _ := datas[i].(int64)
datas[i] = ledis.StrPutInt64(v)
}
}
c.writeArray(datas)
}
return nil