add Sorted.Search.Union for multi search (#12)

xh3b4sd · Oct 8, 2023 · 2f90402 · 2f90402
1 parent b50949e
commit 2f90402
Show file tree

Hide file tree

Showing 4 changed files with 196 additions and 6 deletions.
diff --git a/pkg/conformance/client_single_sorted_test.go b/pkg/conformance/client_single_sorted_test.go
@@ -1607,6 +1607,147 @@ func Test_Client_Single_Sorted_Search_Rando_Cou(t *testing.T) {
 	}
 }
 
+func Test_Client_Single_Sorted_Search_Union(t *testing.T) {
+	var err error
+
+	var cli redigo.Interface
+	{
+		c := client.Config{
+			Kind: client.KindSingle,
+		}
+
+		cli, err = client.New(c)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		err = cli.Purge()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		res, err := cli.Sorted().Search().Union("k1", "k2")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(res) != 0 {
+			t.Fatal("expected", 0, "got", len(res))
+		}
+	}
+
+	{
+		err = cli.Sorted().Create().Index("k1", "v3", 0.3)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cli.Sorted().Create().Index("k1", "v4", 0.4)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cli.Sorted().Create().Index("k1", "v5", 0.5)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cli.Sorted().Create().Index("k1", "v6", 0.6)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		res, err := cli.Sorted().Search().Union("k1")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(res) != 4 {
+			t.Fatal("expected", 4, "got", len(res))
+		}
+		if res[0] != "v3" {
+			t.Fatal("expected", "v3", "got", res[0])
+		}
+		if res[1] != "v4" {
+			t.Fatal("expected", "v4", "got", res[1])
+		}
+		if res[2] != "v5" {
+			t.Fatal("expected", "v5", "got", res[2])
+		}
+		if res[3] != "v6" {
+			t.Fatal("expected", "v6", "got", res[3])
+		}
+	}
+
+	{
+		res, err := cli.Sorted().Search().Union("k1", "k2")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(res) != 4 {
+			t.Fatal("expected", 4, "got", len(res))
+		}
+		if res[0] != "v3" {
+			t.Fatal("expected", "v3", "got", res[0])
+		}
+		if res[1] != "v4" {
+			t.Fatal("expected", "v4", "got", res[1])
+		}
+		if res[2] != "v5" {
+			t.Fatal("expected", "v5", "got", res[2])
+		}
+		if res[3] != "v6" {
+			t.Fatal("expected", "v6", "got", res[3])
+		}
+	}
+
+	{
+		err = cli.Sorted().Create().Index("k2", "v2", 0.2)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cli.Sorted().Create().Index("k2", "v4", 0.4)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cli.Sorted().Create().Index("k2", "v5", 0.5)
+		if err != nil {
+			t.Fatal(err)
+		}
+		err = cli.Sorted().Create().Index("k2", "v7", 0.7)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	{
+		res, err := cli.Sorted().Search().Union("k1", "k2")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(res) != 6 {
+			t.Fatal("expected", 6, "got", len(res))
+		}
+		if res[0] != "v2" {
+			t.Fatal("expected", "v2", "got", res[0])
+		}
+		if res[1] != "v3" {
+			t.Fatal("expected", "v3", "got", res[1])
+		}
+		if res[2] != "v4" {
+			t.Fatal("expected", "v4", "got", res[2])
+		}
+		if res[3] != "v5" {
+			t.Fatal("expected", "v5", "got", res[3])
+		}
+		if res[4] != "v6" {
+			t.Fatal("expected", "v6", "got", res[4])
+		}
+		if res[5] != "v7" {
+			t.Fatal("expected", "v7", "got", res[5])
+		}
+	}
+}
+
 func Test_Client_Single_Sorted_Search_Value(t *testing.T) {
 	var err error
 

diff --git a/pkg/fake/sorted_search.go b/pkg/fake/sorted_search.go
@@ -6,6 +6,7 @@ type SortedSearch struct {
 	FakeOrder func() ([]string, error)
 	FakeRando func() ([]string, error)
 	FakeScore func() ([]string, error)
+	FakeUnion func() ([]string, error)
 }
 
 func (s *SortedSearch) Index(key string, ind string) (string, error) {
@@ -47,3 +48,11 @@ func (s *SortedSearch) Score(key string, lef float64, rig float64) ([]string, er
 
 	return nil, nil
 }
+
+func (s *SortedSearch) Union(key ...string) ([]string, error) {
+	if s.FakeUnion != nil {
+		return s.FakeUnion()
+	}
+
+	return nil, nil
+}
diff --git a/pkg/sorted/interface.go b/pkg/sorted/interface.go
@@ -111,6 +111,10 @@ type Search interface {
 	//     k1       v3 v4 v5 v6
 	//     k2    v2    v4 v5    v7
 	//
+	// For more information about the underlying behaviour see ZINTER.
+	//
+	//     https://redis.io/commands/zinter
+	//
 	Inter(key ...string) ([]string, error)
 
 	// Order returns the values of the sorted set elements stored under key. The
@@ -137,6 +141,20 @@ type Search interface {
 	//     https://redis.io/commands/zrange
 	//
 	Score(key string, lef float64, rig float64) ([]string, error)
+
+	// Union returns the unique values that exist in any of the given keys.
+	// Therefore the returned values represent the union of the given keys. Given
+	// k1 and k2 hold the following values, Union(k1, k2) were to return v2, v3,
+	// v4, v5, v6 and v7.
+	//
+	//     k1       v3 v4 v5 v6
+	//     k2    v2    v4 v5    v7
+	//
+	// For more information about the underlying behaviour see ZUNION.
+	//
+	//     https://redis.io/commands/zunion
+	//
+	Union(key ...string) ([]string, error)
 }
 
 type Update interface {

diff --git a/pkg/sorted/search.go b/pkg/sorted/search.go
@@ -64,10 +64,10 @@ func (s *search) Inter(key ...string) ([]string, error) {
 	var mul []interface{}
 	{
 		mul = append(mul, len(key))
-	}
 
-	for _, x := range key {
-		mul = append(mul, prefix.WithKeys(s.prefix, x))
+		for _, x := range key {
+			mul = append(mul, prefix.WithKeys(s.prefix, x))
+		}
 	}
 
 	res, err := redis.Strings(con.Do("ZINTER", mul...))
@@ -91,9 +91,8 @@ func (s *search) Order(key string, lef int, rig int, sco ...bool) ([]string, err
 
 	var arg []interface{}
 	{
-		arg = append(arg, prefix.WithKeys(s.prefix, key))
-		arg = append(arg, lef)
-		arg = append(arg, rig)
+		arg = append(arg, prefix.WithKeys(s.prefix, key), lef, rig)
+
 		if len(sco) == 1 {
 			arg = append(arg, "WITHSCORES")
 		}
@@ -148,3 +147,26 @@ func (s *search) Score(key string, lef float64, rig float64) ([]string, error) {
 
 	return res, nil
 }
+
+func (s *search) Union(key ...string) ([]string, error) {
+	con := s.pool.Get()
+	defer con.Close()
+
+	var mul []interface{}
+	{
+		mul = append(mul, len(key))
+
+		for _, x := range key {
+			mul = append(mul, prefix.WithKeys(s.prefix, x))
+		}
+
+		mul = append(mul, "AGGREGATE", "MIN")
+	}
+
+	res, err := redis.Strings(con.Do("ZUNION", mul...))
+	if err != nil {
+		return nil, tracer.Mask(err)
+	}
+
+	return res, nil
+}