Add Support for Count-Min Sketch Commands

Author: dengliming

client.go

@@ -131,3 +131,71 @@ func (client *Client) BfExistsMulti(key string, items []string) ([]int64, error)
 	result, err := conn.Do("BF.MEXISTS", args...)
 	return redis.Int64s(result, 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

@@ -129,3 +129,76 @@ func TestClient_BfExistsMulti(t *testing.T) {
 	assert.Equal(t, int64(1), existsResult[1])
 	assert.Equal(t, int64(0), existsResult[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"])
+}