diff --git a/ledis/t_zset.go b/ledis/t_zset.go index 3a2336f..dc028c0 100644 --- a/ledis/t_zset.go +++ b/ledis/t_zset.go @@ -939,3 +939,83 @@ func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, agg func (db *DB) ZScan(key []byte, count int, inclusive bool, match string) ([][]byte, error) { return db.scan(ZSizeType, key, count, inclusive, match) } + +func (db *DB) ZRangeByLex(key []byte, min []byte, max []byte, rangeType uint8, offset int, count int) ([][]byte, error) { + if min == nil { + min = db.zEncodeStartSetKey(key) + } else { + min = db.zEncodeSetKey(key, min) + } + if max == nil { + max = db.zEncodeStopSetKey(key) + } else { + max = db.zEncodeSetKey(key, max) + } + + it := db.bucket.RangeLimitIterator(min, max, rangeType, offset, count) + defer it.Close() + + ay := make([][]byte, 0, 16) + for ; it.Valid(); it.Next() { + if _, m, err := db.zDecodeSetKey(it.Key()); err == nil { + ay = append(ay, m) + } + } + + return ay, nil +} + +func (db *DB) ZRemRangeByLex(key []byte, min []byte, max []byte, rangeType uint8) (int64, error) { + if min == nil { + min = db.zEncodeStartSetKey(key) + } else { + min = db.zEncodeSetKey(key, min) + } + if max == nil { + max = db.zEncodeStopSetKey(key) + } else { + max = db.zEncodeSetKey(key, max) + } + + t := db.zsetBatch + t.Lock() + defer t.Unlock() + + it := db.bucket.RangeIterator(min, max, rangeType) + defer it.Close() + + var n int64 = 0 + for ; it.Valid(); it.Next() { + t.Delete(it.RawKey()) + n++ + } + + if err := t.Commit(); err != nil { + return 0, err + } + + return n, nil +} + +func (db *DB) ZLexCount(key []byte, min []byte, max []byte, rangeType uint8) (int64, error) { + if min == nil { + min = db.zEncodeStartSetKey(key) + } else { + min = db.zEncodeSetKey(key, min) + } + if max == nil { + max = db.zEncodeStopSetKey(key) + } else { + max = db.zEncodeSetKey(key, max) + } + + it := db.bucket.RangeIterator(min, max, rangeType) + defer it.Close() + + var n int64 = 0 + for ; it.Valid(); it.Next() { + n++ + } + + return n, nil +} diff --git a/ledis/t_zset_test.go b/ledis/t_zset_test.go index a1754ed..98f0a81 100644 --- a/ledis/t_zset_test.go +++ b/ledis/t_zset_test.go @@ -2,6 +2,8 @@ package ledis import ( "fmt" + "github.com/siddontang/ledisdb/store" + "reflect" "testing" ) @@ -407,3 +409,59 @@ func TestZScan(t *testing.T) { t.Fatal("invalid value length ", len(v)) } } + +func TestZLex(t *testing.T) { + db := getTestDB() + if _, err := db.zFlush(); err != nil { + t.Fatal(err) + } + + key := []byte("myzset") + if _, err := db.ZAdd(key, ScorePair{0, []byte("a")}, + ScorePair{0, []byte("b")}, + ScorePair{0, []byte("c")}, + ScorePair{0, []byte("d")}, + ScorePair{0, []byte("e")}, + ScorePair{0, []byte("f")}, + ScorePair{0, []byte("g")}); err != nil { + t.Fatal(err) + } + + if ay, err := db.ZRangeByLex(key, nil, []byte("c"), store.RangeClose, 0, -1); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ay, [][]byte{[]byte("a"), []byte("b"), []byte("c")}) { + t.Fatal("must equal a, b, c") + } + + if ay, err := db.ZRangeByLex(key, nil, []byte("c"), store.RangeROpen, 0, -1); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ay, [][]byte{[]byte("a"), []byte("b")}) { + t.Fatal("must equal a, b") + } + + if ay, err := db.ZRangeByLex(key, []byte("aaa"), []byte("g"), store.RangeROpen, 0, -1); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ay, [][]byte{[]byte("b"), + []byte("c"), []byte("d"), []byte("e"), []byte("f")}) { + t.Fatal("must equal b, c, d, e, f", fmt.Sprintf("%q", ay)) + } + + if n, err := db.ZLexCount(key, nil, nil, store.RangeClose); err != nil { + t.Fatal(err) + } else if n != 7 { + t.Fatal(n) + } + + if n, err := db.ZRemRangeByLex(key, []byte("aaa"), []byte("g"), store.RangeROpen); err != nil { + t.Fatal(err) + } else if n != 5 { + t.Fatal(n) + } + + if n, err := db.ZLexCount(key, nil, nil, store.RangeClose); err != nil { + t.Fatal(err) + } else if n != 2 { + t.Fatal(n) + } + +} diff --git a/server/cmd_zset.go b/server/cmd_zset.go index 4763683..3c5abcb 100644 --- a/server/cmd_zset.go +++ b/server/cmd_zset.go @@ -5,6 +5,7 @@ import ( "github.com/siddontang/go/hack" "github.com/siddontang/go/num" "github.com/siddontang/ledisdb/ledis" + "github.com/siddontang/ledisdb/store" "math" "strconv" "strings" @@ -661,6 +662,128 @@ func zxscanCommand(c *client) error { return nil } +func zparseMemberRange(minBuf []byte, maxBuf []byte) (min []byte, max []byte, rangeType uint8, err error) { + rangeType = store.RangeClose + if strings.ToLower(hack.String(minBuf)) == "-" { + min = nil + } else { + if len(minBuf) == 0 { + err = ErrCmdParams + return + } + + if minBuf[0] == '(' { + rangeType |= store.RangeLOpen + min = minBuf[1:] + } else if minBuf[0] == '[' { + min = minBuf[1:] + } else { + err = ErrCmdParams + return + } + } + + if strings.ToLower(hack.String(maxBuf)) == "+" { + max = nil + } else { + if len(maxBuf) == 0 { + err = ErrCmdParams + return + } + if maxBuf[0] == '(' { + rangeType |= store.RangeROpen + max = maxBuf[1:] + } else if maxBuf[0] == '[' { + max = maxBuf[1:] + } else { + err = ErrCmdParams + return + } + } + + return +} + +func zrangebylexCommand(c *client) error { + args := c.args + if len(args) != 3 && len(args) != 6 { + return ErrCmdParams + } + + min, max, rangeType, err := zparseMemberRange(args[1], args[2]) + if err != nil { + return err + } + + var offset int = 0 + var count int = -1 + + if len(args) == 6 { + if strings.ToLower(hack.String(args[3])) != "limit" { + return ErrSyntax + } + + if offset, err = strconv.Atoi(hack.String(args[4])); err != nil { + return ErrValue + } + + if count, err = strconv.Atoi(hack.String(args[5])); err != nil { + return ErrValue + } + } + + key := args[0] + if ay, err := c.db.ZRangeByLex(key, min, max, rangeType, offset, count); err != nil { + return err + } else { + c.resp.writeSliceArray(ay) + } + + return nil +} + +func zremrangebylexCommand(c *client) error { + args := c.args + if len(args) != 3 { + return ErrCmdParams + } + + min, max, rangeType, err := zparseMemberRange(args[1], args[2]) + if err != nil { + return err + } + + key := args[0] + if n, err := c.db.ZRemRangeByLex(key, min, max, rangeType); err != nil { + return err + } else { + c.resp.writeInteger(n) + } + + return nil +} + +func zlexcountCommand(c *client) error { + args := c.args + if len(args) != 3 { + return ErrCmdParams + } + + min, max, rangeType, err := zparseMemberRange(args[1], args[2]) + if err != nil { + return err + } + + key := args[0] + if n, err := c.db.ZLexCount(key, min, max, rangeType); err != nil { + return err + } else { + c.resp.writeInteger(n) + } + + return nil +} + func init() { register("zadd", zaddCommand) register("zcard", zcardCommand) @@ -680,6 +803,10 @@ func init() { register("zunionstore", zunionstoreCommand) register("zinterstore", zinterstoreCommand) + register("zrangebylex", zrangebylexCommand) + register("zremrangebylex", zremrangebylexCommand) + register("zlexcount", zlexcountCommand) + //ledisdb special command register("zclear", zclearCommand) diff --git a/server/cmd_zset_test.go b/server/cmd_zset_test.go index 8c74bdc..59411c5 100644 --- a/server/cmd_zset_test.go +++ b/server/cmd_zset_test.go @@ -3,6 +3,7 @@ package server import ( "fmt" "github.com/siddontang/ledisdb/client/go/ledis" + "reflect" "strconv" "testing" ) @@ -737,3 +738,51 @@ func TestZInterStore(t *testing.T) { } } } + +func TestZSetLex(t *testing.T) { + c := getTestConn() + defer c.Close() + + key := []byte("myzlexset") + if _, err := c.Do("zadd", key, + 0, "a", 0, "b", 0, "c", 0, "d", 0, "e", 0, "f", 0, "g"); err != nil { + t.Fatal(err) + } + + if ay, err := ledis.Strings(c.Do("zrangebylex", key, "-", "[c")); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ay, []string{"a", "b", "c"}) { + t.Fatal("must equal") + } + + if ay, err := ledis.Strings(c.Do("zrangebylex", key, "-", "(c")); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ay, []string{"a", "b"}) { + t.Fatal("must equal") + } + + if ay, err := ledis.Strings(c.Do("zrangebylex", key, "[aaa", "(g")); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ay, []string{"b", "c", "d", "e", "f"}) { + t.Fatal("must equal") + } + + if n, err := ledis.Int64(c.Do("zlexcount", key, "-", "(c")); err != nil { + t.Fatal(err) + } else if n != 2 { + t.Fatal(n) + } + + if n, err := ledis.Int64(c.Do("zremrangebylex", key, "[aaa", "(g")); err != nil { + t.Fatal(err) + } else if n != 5 { + t.Fatal(n) + } + + if n, err := ledis.Int64(c.Do("zlexcount", key, "-", "+")); err != nil { + t.Fatal(err) + } else if n != 2 { + t.Fatal(n) + } + +}