Merge pull request #37 from RedisBloom/support_TDigest

Add support for T-Digest sketch to the client

Author: AvitalFineRedis

client.go

@@ -18,6 +18,52 @@ type Client struct {
 	Name string
 }
 
+// TDigestInfo is a struct that represents T-Digest properties
+type TDigestInfo struct {
+	compression       int64
+	capacity          int64
+	mergedNodes       int64
+	unmergedNodes     int64
+	mergedWeight      float64
+	unmergedWeight    float64
+	totalCompressions int64
+}
+
+// Compression - returns the compression of TDigestInfo instance
+func (info *TDigestInfo) Compression() int64 {
+	return info.compression
+}
+
+// Capacity - returns the capacity of TDigestInfo instance
+func (info *TDigestInfo) Capacity() int64 {
+	return info.capacity
+}
+
+// MergedNodes - returns the merged nodes of TDigestInfo instance
+func (info *TDigestInfo) MergedNodes() int64 {
+	return info.mergedNodes
+}
+
+// UnmergedNodes -  returns the unmerged nodes of TDigestInfo instance
+func (info *TDigestInfo) UnmergedNodes() int64 {
+	return info.unmergedNodes
+}
+
+// MergedWeight - returns the merged weight of TDigestInfo instance
+func (info *TDigestInfo) MergedWeight() float64 {
+	return info.mergedWeight
+}
+
+// UnmergedWeight - returns the unmerged weight of TDigestInfo instance
+func (info *TDigestInfo) UnmergedWeight() float64 {
+	return info.unmergedWeight
+}
+
+// TotalCompressions - returns the total compressions of TDigestInfo instance
+func (info *TDigestInfo) TotalCompressions() int64 {
+	return info.totalCompressions
+}
+
 // NewClient creates a new client connecting to the redis host, and using the given name as key prefix.
 // Addr can be a single host:port pair, or a comma separated list of host:port,host:port...
 // In the case of multiple hosts we create a multi-pool and select connections at random
@@ -443,6 +489,76 @@ func (client *Client) CfInfo(key string) (map[string]int64, error) {
 	return ParseInfoReply(redis.Values(conn.Do("CF.INFO", key)))
 }
 
+// TdCreate - Allocate the memory and initialize the t-digest
+func (client *Client) TdCreate(key string, compression int64) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.String(conn.Do("TDIGEST.CREATE", key, compression))
+}
+
+// TdReset - Reset the sketch to zero - empty out the sketch and re-initialize it
+func (client *Client) TdReset(key string) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.String(conn.Do("TDIGEST.RESET", key))
+}
+
+// TdAdd - Adds one or more samples to a sketch
+func (client *Client) TdAdd(key string, samples map[float64]float64) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	args := redis.Args{key}
+	for k, v := range samples {
+		args = args.Add(k, v)
+	}
+	reply, err := conn.Do("TDIGEST.ADD", args...)
+	return redis.String(reply, err)
+}
+
+// TdMerge - Merges all of the values from 'from' to 'this' sketch
+func (client *Client) TdMerge(toKey string, fromKey string) (string, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.String(conn.Do("TDIGEST.MERGE", toKey, fromKey))
+}
+
+// TdMin - Get minimum value from the sketch. Will return DBL_MAX if the sketch is empty
+func (client *Client) TdMin(key string) (float64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.Float64(conn.Do("TDIGEST.MIN", key))
+}
+
+// TdMax - Get maximum value from the sketch. Will return DBL_MIN if the sketch is empty
+func (client *Client) TdMax(key string) (float64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.Float64(conn.Do("TDIGEST.MAX", key))
+}
+
+// TdQuantile - Returns an estimate of the cutoff such that a specified fraction of the data added
+// to this TDigest would be less than or equal to the cutoff
+func (client *Client) TdQuantile(key string, quantile float64) (float64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.Float64(conn.Do("TDIGEST.QUANTILE", key, quantile))
+}
+
+// TdCdf - Returns the fraction of all points added which are <= value
+func (client *Client) TdCdf(key string, value float64) (float64, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return redis.Float64(conn.Do("TDIGEST.CDF", key, value))
+}
+
+// TdInfo - Returns compression, capacity, total merged and unmerged nodes, the total
+// compressions made up to date on that key, and merged and unmerged weight.
+func (client *Client) TdInfo(key string) (TDigestInfo, error) {
+	conn := client.Pool.Get()
+	defer conn.Close()
+	return ParseTDigestInfo(redis.Values(conn.Do("TDIGEST.INFO", key)))
+}
+
 func ParseInfoReply(values []interface{}, err error) (map[string]int64, error) {
 	if err != nil {
 		return nil, err
@@ -456,3 +572,38 @@ func ParseInfoReply(values []interface{}, err error) (map[string]int64, error) {
 	}
 	return m, err
 }
+
+func ParseTDigestInfo(result interface{}, err error) (info TDigestInfo, outErr error) {
+	values, outErr := redis.Values(result, err)
+	if outErr != nil {
+		return TDigestInfo{}, err
+	}
+	if len(values)%2 != 0 {
+		return TDigestInfo{}, errors.New("ParseInfo expects even number of values result")
+	}
+	var key string
+	for i := 0; i < len(values); i += 2 {
+		key, outErr = redis.String(values[i], nil)
+		switch key {
+		case "Compression":
+			info.compression, outErr = redis.Int64(values[i+1], nil)
+		case "Capacity":
+			info.capacity, outErr = redis.Int64(values[i+1], nil)
+		case "Merged nodes":
+			info.mergedNodes, outErr = redis.Int64(values[i+1], nil)
+		case "Unmerged nodes":
+			info.unmergedNodes, outErr = redis.Int64(values[i+1], nil)
+		case "Merged weight":
+			info.mergedWeight, outErr = redis.Float64(values[i+1], nil)
+		case "Unmerged weight":
+			info.unmergedWeight, outErr = redis.Float64(values[i+1], nil)
+		case "Total compressions":
+			info.totalCompressions, outErr = redis.Int64(values[i+1], nil)
+		}
+		if outErr != nil {
+			return TDigestInfo{}, outErr
+		}
+	}
+
+	return info, nil
+}

client_test.go

@@ -527,5 +527,128 @@ func TestClient_BfScanDump(t *testing.T) {
 	assert.Nil(t, err)
 	_, _, err = client.BfScanDump(notBfKey, 0)
 	assert.Equal(t, err.Error(), "WRONGTYPE Operation against a key holding the wrong kind of value")
+}
+
+func TestClient_TdReset(t *testing.T) {
+	client.FlushAll()
+	key := "test_td"
+	ret, err := client.TdCreate(key, 100)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	ret, err = client.TdReset(key)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	samples := map[float64]float64{1.0: 1.0, 2.0: 2.0}
+	ret, err = client.TdAdd(key, samples)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	ret, err = client.TdReset(key)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	info, err := client.TdInfo(key)
+	assert.Nil(t, err)
+	assert.Equal(t, 0.0, info.UnmergedWeight())
+	assert.Equal(t, int64(0), info.TotalCompressions())
+	assert.Equal(t, int64(100), info.Compression())
+	assert.Equal(t, int64(610), info.Capacity())
+}
 
+func TestClient_TdMerge(t *testing.T) {
+	key1 := "toKey"
+	key2 := "fromKey"
+	ret, err := client.TdCreate(key1, 10)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	ret, err = client.TdCreate(key2, 10)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	//Add values
+	samples1 := map[float64]float64{1.0: 1.0, 2.0: 2.0}
+	samples2 := map[float64]float64{3.0: 3.0, 4.0: 4.0}
+	ret, err = client.TdAdd(key1, samples1)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+	ret, err = client.TdAdd(key2, samples2)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	//Merge
+	ret, err = client.TdMerge(key1, key2)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	// we should now have 10 weight on to-histogram
+	info, err := client.TdInfo(key1)
+	assert.Nil(t, err)
+	assert.Equal(t, 10.0, info.UnmergedWeight()+info.MergedWeight())
+	assert.Equal(t, int64(2), info.UnmergedNodes())
+	assert.Equal(t, int64(2), info.MergedNodes())
+}
+
+func TestClient_TdMinMax(t *testing.T) {
+	client.FlushAll()
+	key := "test_td"
+	ret, err := client.TdCreate(key, 10)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	samples := map[float64]float64{1.0: 1.0, 2.0: 2.0, 3.0: 3.0}
+	ret, err = client.TdAdd(key, samples)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	ans, err := client.TdMin(key)
+	assert.Nil(t, err)
+	assert.Equal(t, 1.0, ans)
+
+	ans, err = client.TdMax(key)
+	assert.Nil(t, err)
+	assert.Equal(t, 3.0, ans)
+}
+
+func TestClient_TdQuantile(t *testing.T) {
+	client.FlushAll()
+	key := "test_td"
+	ret, err := client.TdCreate(key, 10)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	samples := map[float64]float64{1.0: 1.0, 2.0: 1.0, 3.0: 1.0}
+	ret, err = client.TdAdd(key, samples)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	ans, err := client.TdQuantile(key, 1.0)
+	assert.Nil(t, err)
+	assert.Equal(t, 3.0, ans)
+
+	ans, err = client.TdQuantile(key, 0.0)
+	assert.Nil(t, err)
+	assert.Equal(t, 1.0, ans)
+}
+
+func TestClient_TdCdf(t *testing.T) {
+	client.FlushAll()
+	key := "test_td"
+	ret, err := client.TdCreate(key, 10)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	samples := map[float64]float64{1.0: 1.0, 2.0: 1.0, 3.0: 1.0}
+	ret, err = client.TdAdd(key, samples)
+	assert.Nil(t, err)
+	assert.Equal(t, "OK", ret)
+
+	ans, err := client.TdCdf(key, 10.0)
+	assert.Nil(t, err)
+	assert.Equal(t, 1.0, ans)
+
+	ans, err = client.TdCdf(key, 0.0)
+	assert.Nil(t, err)
+	assert.Equal(t, 0.0, ans)
 }