Merge pull request #16 from dengliming/f_cms_command

Add Support for Count-Min Sketch Commands

Author: filipecosta90

README.md

@@ -61,10 +61,10 @@ func main() {
 | :---          |  ----: |
 | [BF.RESERVE](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfreserve) | [Reserve](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.Reserve) |
 | [BF.ADD](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfadd) | [Add](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.Add) |
-| [BF.MADD](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfmadd) | N/A |
+| [BF.MADD](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfmadd) | [BfAddMulti](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.BfAddMulti)  |
 | [BF.INSERT](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfinsert) | N/A |
 | [BF.EXISTS](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfexists) | [Exists](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.Exists) |
-| [BF.MEXISTS](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfmexists) | N/A |
+| [BF.MEXISTS](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfmexists) | [BfExistsMulti](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.BfExistsMulti) |
 | [BF.SCANDUMP](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfscandump) | N/A |
 | [BF.LOADCHUNK](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfloadchunk) | N/A |
 | [BF.INFO](https://oss.redislabs.com/redisbloom/Bloom_Commands/#bfinfo) | [Info](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.Info) |
@@ -89,12 +89,12 @@ func main() {
 
 | Command | Recommended API and godoc  |
 | :---          |  ----: |
-| [CMS.INITBYDIM](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsinitbydim) | N/A |
-| [CMS.INITBYPROB](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsinitbyprob) |  N/A |
-| [CMS.INCRBY](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsincrby) |  N/A |
-| [CMS.QUERY](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsquery) |  N/A |
-| [CMS.MERGE](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsmerge) |  N/A |
-| [CMS.INFO](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsinfo) |  N/A |
+| [CMS.INITBYDIM](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsinitbydim) | [CmsInitByDim](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.CmsInitByDim) |
+| [CMS.INITBYPROB](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsinitbyprob) |  [CmsInitByProb](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.CmsInitByProb) |
+| [CMS.INCRBY](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsincrby) |  [CmsIncrBy](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.CmsIncrBy) |
+| [CMS.QUERY](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsquery) | [CmsQuery](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.CmsQuery) |
+| [CMS.MERGE](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsmerge) |  [CmsMerge](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.CmsMerge) |
+| [CMS.INFO](https://oss.redislabs.com/redisbloom/CountMinSketch_Commands/#cmsinfo) |  [CmsInfo](https://godoc.org/github.com/RedisBloom/redisbloom-go#Client.CmsInfo) |
 
 ### TopK Filter
 

client.go

@@ -216,3 +216,71 @@ func (client *Client) TopkIncrBy(key string, itemIncrements map[string]int64) ([
 	reply, err := conn.Do("TOPK.INCRBY", args...)
 	return redis.Strings(reply, err)
 }
+
+// Initializes a Count-Min Sketch to dimensions specified by user.
+func (client *Client) CmsInitByDim(key string, width int64, depth int64) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	result, err := conn.Do("CMS.INITBYDIM", key, width, depth)
+	return redis.String(result, err)
+}
+
+// Initializes a Count-Min Sketch to accommodate requested capacity.
+func (client *Client) CmsInitByProb(key string, error float64, probability float64) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	result, err := conn.Do("CMS.INITBYPROB", key, error, probability)
+	return redis.String(result, err)
+}
+
+// Increases the count of item by increment. Multiple items can be increased with one call.
+func (client *Client) CmsIncrBy(key string, itemIncrements map[string]int64) ([]int64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}
+	for k, v := range itemIncrements {
+		args = args.Add(k, v)
+	}
+	result, err := conn.Do("CMS.INCRBY", args...)
+	return redis.Int64s(result, err)
+}
+
+// Returns count for item.
+func (client *Client) CmsQuery(key string, items []string) ([]int64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}.AddFlat(items)
+	result, err := conn.Do("CMS.QUERY", args...)
+	return redis.Int64s(result, err)
+}
+
+// Merges several sketches into one sketch.
+func (client *Client) CmsMerge(key string, srcs []string, weights []string) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}.Add(len(srcs)).AddFlat(srcs)
+	if weights != nil && len(weights) > 0 {
+		args = args.Add("WEIGHTS").AddFlat(weights)
+	}
+	return redis.String(conn.Do("CMS.MERGE", args...))
+}
+
+// Returns width, depth and total count of the sketch.
+func (client *Client) CmsInfo(key string) (map[string]int64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	reply, err := conn.Do("CMS.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]int64, len(values)/2)
+	for i := 0; i < len(values); i += 2 {
+		m[values[i].(string)] = values[i+1].(int64)
+	}
+	return m, err
+}

client_test.go

@@ -207,3 +207,76 @@ func TestClient_TopkIncrBy(t *testing.T) {
 	assert.Equal(t, 3, len(rets))
 	assert.Equal(t, "", rets[2])
 }
+
+func TestClient_CmsInitByDim(t *testing.T) {
+	client.FlushAll()
+	ret, err := client.CmsInitByDim("test_cms_initbydim", 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+}
+
+func TestClient_CmsInitByProb(t *testing.T) {
+	client.FlushAll()
+	ret, err := client.CmsInitByProb("test_cms_initbyprob", 0.01, 0.01)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+}
+
+func TestClient_CmsIncrBy(t *testing.T) {
+	client.FlushAll()
+	key := "test_cms_incrby"
+	ret, err := client.CmsInitByDim(key, 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	results, err := client.CmsIncrBy(key, map[string]int64{"foo": 5})
+	assert.Nil(t, err)
+	assert.NotNil(t, results)
+	assert.Equal(t, int64(5), results[0])
+}
+
+func TestClient_CmsQuery(t *testing.T) {
+	client.FlushAll()
+	key := "test_cms_query"
+	ret, err := client.CmsInitByDim(key, 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	results, err := client.CmsQuery(key, []string{"notexist"})
+	assert.Nil(t, err)
+	assert.NotNil(t, 0, results[0])
+	_, err = client.CmsIncrBy(key, map[string]int64{"foo": 5})
+	assert.Nil(t, err)
+	results, err = client.CmsQuery(key, []string{"foo"})
+	assert.Nil(t, err)
+	assert.Equal(t, int64(5), results[0])
+}
+
+func TestClient_CmsMerge(t *testing.T) {
+	client.FlushAll()
+	ret, err := client.CmsInitByDim("A", 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	ret, err = client.CmsInitByDim("B", 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	ret, err = client.CmsInitByDim("C", 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	client.CmsIncrBy("A", map[string]int64{"foo": 5, "bar": 3, "baz": 9})
+	client.CmsIncrBy("B", map[string]int64{"foo": 2, "bar": 3, "baz": 1})
+	client.CmsMerge("C", []string{"A", "B"}, nil)
+	results, err := client.CmsQuery("C", []string{"foo", "bar", "baz"})
+	assert.Equal(t, []int64{7, 6, 10}, results)
+}
+
+func TestClient_CmsInfo(t *testing.T) {
+	client.FlushAll()
+	key := "test_cms_info"
+	ret, err := client.CmsInitByDim(key, 1000, 5)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	info, err := client.CmsInfo(key)
+	assert.Nil(t, err)
+	assert.Equal(t, int64(1000), info["width"])
+	assert.Equal(t, int64(5), info["depth"])
+	assert.Equal(t, int64(0), info["count"])
+}