mirror
/
ledisdb
mirror of https://github.com/ledisdb/ledisdb.git
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'develop' into scan-feature

This commit is contained in:
wenyekui 2014-08-14 16:14:40 +08:00
parent 54ed8e0d38 34f1795a5f
commit 647790674b
2 changed files with 50 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
ledis/t_zset.go
View File

@ -939,7 +939,6 @@ func (db *DB) ZUnionStore(destKey []byte, srcKeys [][]byte, weights []int64, agg
func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, aggregate byte) (int64, error) { func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, aggregate byte) (int64, error) {
var destMap = map[string]int64{}
aggregateFunc := getAggregateFunc(aggregate) aggregateFunc := getAggregateFunc(aggregate)
if aggregateFunc == nil { if aggregateFunc == nil {
return 0, errInvalidAggregate return 0, errInvalidAggregate
@ -958,20 +957,27 @@ func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, agg
} }
} }
var keptMembers [][]byte var destMap = map[string]int64{}
for i, key := range srcKeys { scorePairs, err := db.ZRange(srcKeys[0], 0, -1)
scorePairs, err := db.ZRange(key, 0, -1)
if err != nil { if err != nil {
return 0, err return 0, err
} }
for _, pair := range scorePairs { for _, pair := range scorePairs {
if score, ok := destMap[String(pair.Member)]; !ok { destMap[String(pair.Member)] = pair.Score * weights[0]
destMap[String(pair.Member)] = pair.Score * weights[i] }
} else {
keptMembers = append(keptMembers, pair.Member) for i, key := range srcKeys[1:] {
destMap[String(pair.Member)] = aggregateFunc(score, pair.Score*weights[i]) scorePairs, err := db.ZRange(key, 0, -1)
if err != nil {
return 0, err
}
tmpMap := map[string]int64{}
for _, pair := range scorePairs {
if score, ok := destMap[String(pair.Member)]; ok {
tmpMap[String(pair.Member)] = aggregateFunc(score, pair.Score*weights[i+1])
} }
} }
destMap = tmpMap
} }
t := db.zsetTx t := db.zsetTx
@ -981,13 +987,12 @@ func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, agg
db.zDelete(t, destKey) db.zDelete(t, destKey)
var num int64 = 0 var num int64 = 0
for _, member := range keptMembers { for member, score := range destMap {
score := destMap[String(member)] if err := checkZSetKMSize(destKey, []byte(member)); err != nil {
if err := checkZSetKMSize(destKey, member); err != nil {
return 0, err return 0, err
} }
if n, err := db.zSetItem(t, destKey, score, member); err != nil { if n, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
return 0, err return 0, err
} else if n == 0 { } else if n == 0 {
//add new //add new
@ -1002,7 +1007,7 @@ func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, agg
if err := t.Commit(); err != nil { if err := t.Commit(); err != nil {
return 0, err return 0, err
} }
return int64(len(keptMembers)), nil return int64(len(destMap)), nil
} }
func (db *DB) ZScan(key []byte, count int, inclusive bool) ([][]byte, error) { func (db *DB) ZScan(key []byte, count int, inclusive bool) ([][]byte, error) {

32
server/cmd_zset_test.go
View File

@ -704,4 +704,36 @@ func TestZInterStore(t *testing.T) {
t.Fatal("invalid value ", n) t.Fatal("invalid value ", n)
} }
} }
if _, err := c.Do("zadd", "k3", "3", "three"); err != nil {
t.Fatal(err.Error())
}
if n, err := ledis.Int64(c.Do("zinterstore", "out", "3", "k1", "k2", "k3", "aggregate", "sum")); err != nil {
t.Fatal(err.Error())
} else {
if n != 0 {
t.Fatal("invalid value ", n)
}
}
if _, err := c.Do("zadd", "k3", "3", "two"); err != nil {
t.Fatal(err.Error())
}
if n, err := ledis.Int64(c.Do("zinterstore", "out", "3", "k1", "k2", "k3", "aggregate", "sum", "weights", "3", "2", "2")); err != nil {
t.Fatal(err.Error())
} else {
if n != 1 {
t.Fatal("invalid value ", n)
}
}
if n, err := ledis.Int64(c.Do("zscore", "out", "two")); err != nil {
t.Fatal(err.Error())
} else {
if n != 14 {
t.Fatal("invalid value ", n)
}
}
} }