Merge pull request #12 from dengliming/master

Add Support for TopK Filter Commands

Author: gkorland

README.md

@@ -100,13 +100,13 @@ func main() {
 
 | Command | Recommended API and godoc  |
 | :---          |  ----: |
-| [TOPK.RESERVE](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkreserve) | N/A |
-| [TOPK.ADD](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkadd) |  N/A |
-| [TOPK.INCRBY](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkincrby) |  N/A |
-| [TOPK.QUERY](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkquery) |  N/A |
-| [TOPK.COUNT](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkcount) |  N/A |
-| [TOPK.LIST](https://oss.redislabs.com/redisbloom/TopK_Commands/#topklist) |  N/A |
-| [TOPK.INFO](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkinfo) |  N/A |
+| [TOPK.RESERVE](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkreserve) |  [TopkReserve](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkReserve)  |
+| [TOPK.ADD](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkadd) |   [TopkAdd](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkAdd)  |
+| [TOPK.INCRBY](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkincrby) |  [TopkIncrby](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkIncrby)  |
+| [TOPK.QUERY](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkquery) |   [TopkQuery](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkQuery)  |
+| [TOPK.COUNT](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkcount) |   [TopkCount](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkCount)  |
+| [TOPK.LIST](https://oss.redislabs.com/redisbloom/TopK_Commands/#topklist) |   [TopkList](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkList)  |
+| [TOPK.INFO](https://oss.redislabs.com/redisbloom/TopK_Commands/#topkinfo) |   [TopkInfo](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.TopkInfo)  |
 
 
 ## License

client.go

@@ -2,6 +2,7 @@ package redis_bloom_go
 
 import (
 	"errors"
+	"fmt"
 	"github.com/gomodule/redigo/redis"
 	"strconv"
 	"strings"
@@ -131,3 +132,87 @@ func (client *Client) BfExistsMulti(key string, items []string) ([]int64, error)
 	result, err := conn.Do("BF.MEXISTS", args...)
 	return redis.Int64s(result, err)
 }
+
+// Initializes a TopK with specified parameters.
+func (client *Client) TopkReserve(key string, topk int64, width int64, depth int64, decay float64) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	result, err := conn.Do("TOPK.RESERVE", key, topk, width, depth, strconv.FormatFloat(decay, 'g', 16, 64))
+	return redis.String(result, err)
+}
+
+// Adds an item to the data structure.
+func (client *Client) TopkAdd(key string, items []string) ([]string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}.AddFlat(items)
+	result, err := conn.Do("TOPK.ADD", args...)
+	return redis.Strings(result, err)
+}
+
+// Returns count for an item.
+func (client *Client) TopkCount(key string, items []string) ([]string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}.AddFlat(items)
+	result, err := conn.Do("TOPK.COUNT", args...)
+	return redis.Strings(result, err)
+}
+
+// Checks whether an item is one of Top-K items.
+func (client *Client) TopkQuery(key string, items []string) ([]int64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}.AddFlat(items)
+	result, err := conn.Do("TOPK.QUERY", args...)
+	return redis.Int64s(result, err)
+}
+
+// Return full list of items in Top K list.
+func (client *Client) TopkList(key string) ([]string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	result, err := conn.Do("TOPK.LIST", key)
+	return redis.Strings(result, err)
+}
+
+// Returns number of required items (k), width, depth and decay values.
+func (client *Client) TopkInfo(key string) (map[string]string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	reply, err := conn.Do("TOPK.INFO", key)
+	values, err := redis.Values(reply, err)
+	if err != nil {
+		return nil, err
+	}
+	if len(values)%2 != 0 {
+		return nil, errors.New("expects even number of values result")
+	}
+
+	m := make(map[string]string, len(values)/2)
+	for i := 0; i < len(values); i += 2 {
+		k := values[i].(string)
+		switch v := values[i+1].(type) {
+		case []byte:
+			m[k] = string(values[i+1].([]byte))
+			break
+		case int64:
+			m[k] = strconv.FormatInt(values[i+1].(int64), 10)
+		default:
+			return nil, fmt.Errorf("unexpected element type for (Ints,String), got type %T", v)
+		}
+	}
+	return m, err
+}
+
+// Increase the score of an item in the data structure by increment.
+func (client *Client) TopkIncrBy(key string, itemIncrements map[string]int64) ([]string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}
+	for k, v := range itemIncrements {
+		args = args.Add(k, v)
+	}
+	reply, err := conn.Do("TOPK.INCRBY", args...)
+	return redis.Strings(reply, err)
+}

client_test.go

@@ -129,3 +129,81 @@ func TestClient_BfExistsMulti(t *testing.T) {
 	assert.Equal(t, int64(1), existsResult[1])
 	assert.Equal(t, int64(0), existsResult[2])
 }
+
+func TestClient_TopkReserve(t *testing.T) {
+	client.FlushAll()
+	ret, err := client.TopkReserve("test_topk_reserve", 10, 2000, 7, 0.925)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+}
+
+func TestClient_TopkAdd(t *testing.T) {
+	client.FlushAll()
+	key := "test_topk_add"
+	ret, err := client.TopkReserve(key, 10, 2000, 7, 0.925)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	rets, err := client.TopkAdd(key, []string{"test", "test1", "test3"})
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(rets))
+}
+
+func TestClient_TopkQuery(t *testing.T) {
+	client.FlushAll()
+	key := "test_topk_query"
+	ret, err := client.TopkReserve(key, 10, 2000, 7, 0.925)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	rets, err := client.TopkAdd(key, []string{"test"})
+	assert.Nil(t, err)
+	assert.NotNil(t, rets)
+	queryRet, err := client.TopkQuery(key, []string{"test", "nonexist"})
+	assert.Nil(t, err)
+	assert.Equal(t, 2, len(queryRet))
+	assert.Equal(t, int64(1), queryRet[0])
+	assert.Equal(t, int64(0), queryRet[1])
+
+	key1 := "test_topk_list"
+	ret, err = client.TopkReserve(key1, 3, 50, 3, 0.9)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	client.TopkAdd(key1, []string{"A", "B", "C", "D", "E", "A", "A", "B", "C",
+		"G", "D", "B", "D", "A", "E", "E"})
+	keys, err := client.TopkList(key1)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(keys))
+	assert.Equal(t, []string{"D", "A", "B"}, keys)
+}
+
+func TestClient_TopkInfo(t *testing.T) {
+	client.FlushAll()
+	key := "test_topk_info"
+	ret, err := client.TopkReserve(key, 10, 2000, 7, 0.925)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	info, err := client.TopkInfo(key)
+	assert.Equal(t, "10", info["k"])
+	assert.Equal(t, "2000", info["width"])
+	assert.Equal(t, "7", info["depth"])
+
+	info, err = client.TopkInfo("notexists")
+	assert.NotNil(t, err)
+}
+
+func TestClient_TopkIncrBy(t *testing.T) {
+	client.FlushAll()
+	key := "test_topk_incrby"
+	ret, err := client.TopkReserve(key, 50, 2000, 7, 0.925)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	rets, err := client.TopkAdd(key, []string{"foo", "bar", "42"})
+	assert.Nil(t, err)
+	assert.NotNil(t, rets)
+
+	rets, err = client.TopkIncrBy(key, map[string]int64{"foo": 3, "bar": 2, "42": 30})
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(rets))
+	assert.Equal(t, "", rets[2])
+}