Added ZRangeStore

Signed-off-by: Edward Liang <[email protected]>
valkey-io · Feb 7, 2025 · 0b18b51 · 0b18b51
1 parent e651fdb
commit 0b18b51
Show file tree

Hide file tree

Showing 3 changed files with 201 additions and 0 deletions.
diff --git a/go/api/base_client.go b/go/api/base_client.go
@@ -4711,6 +4711,63 @@ func (client *baseClient) ZRangeWithScores(
 	return handleStringDoubleMapResponse(result)
 }
 
+// Stores a specified range of elements from the sorted set at `key`, into a new
+// sorted set at `destination`. If `destination` doesn't exist, a new sorted
+// set is created; if it exists, it's overwritten.
+//
+// See [valkey.io] for more details.
+//
+// Parameters:
+//
+//	destination - The key for the destination sorted set.
+//	key - The key of the source sorted set.
+//	rangeQuery - The range query object representing the type of range query to perform.
+//	 - For range queries by index (rank), use [RangeByIndex].
+//	 - For range queries by lexicographical order, use [RangeByLex].
+//	 - For range queries by score, use [RangeByScore].
+//
+// Return value:
+//
+//	The number of elements in the resulting sorted set.
+//
+// For example:
+//
+//	client.ZAdd("my_sorted_set", map[string]float64{"a": 1.0, "b": 2.0, "c": 3.0})
+//
+//	// Retrieve and store all members of a sorted set in ascending order
+//	res1, err := client.ZRangeStore("my_dest", "my_sorted_set", options.NewRangeByIndexQuery(0, -1))
+//
+//	// Retrieve members within a score range in descending order
+//	query := options.NewRangeByScoreQuery(
+//	    options.NewScoreBoundary(3, false),
+//		options.NewInfiniteScoreBoundary(options.NegativeInfinity)).
+//		SetReverse()
+//	res2, err := client.ZRange("my_dest", query)
+//	fmt.Println(res1)
+//	fmt.Println(res2)
+//
+//	// Output:
+//	// 3
+//	// [b a]
+//
+// [valkey.io]: https://valkey.io/commands/zrangestore/
+func (client *baseClient) ZRangeStore(
+	destination string,
+	key string,
+	rangeQuery options.ZRangeQuery,
+) (int64, error) {
+	args := make([]string, 0, 10)
+	args = append(args, destination)
+	args = append(args, key)
+	args = append(args, rangeQuery.ToArgs()...)
+	result, err := client.executeCommand(C.ZRangeStore, args)
+	if err != nil {
+		return defaultIntResponse, err
+	}
+
+	return handleIntResponse(result)
+}
+
 // Removes the existing timeout on key, turning the key from volatile
 // (a key with an expire set) to persistent (a key that will never expire as no timeout is associated).
 //

diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go
@@ -49,6 +49,8 @@ type SortedSetCommands interface {
 
 	ZRangeWithScores(key string, rangeQuery options.ZRangeQueryWithScores) (map[string]float64, error)
 
+	ZRangeStore(destination string, key string, rangeQuery options.ZRangeQuery) (int64, error)
+
 	ZRank(key string, member string) (Result[int64], error)
 
 	ZRankWithScore(key string, member string) (Result[int64], Result[float64], error)

diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go
@@ -5112,6 +5112,148 @@ func (suite *GlideTestSuite) TestZRangeWithScores() {
 	})
 }
 
+func (suite *GlideTestSuite) TestZRangeStore() {
+	suite.runWithDefaultClients(func(client api.BaseClient) {
+		t := suite.T()
+		key := "{key}" + uuid.New().String()
+		dest := "{key}" + uuid.New().String()
+		memberScoreMap := map[string]float64{
+			"a": 1.0,
+			"b": 2.0,
+			"c": 3.0,
+		}
+		_, err := client.ZAdd(key, memberScoreMap)
+		assert.NoError(t, err)
+		// index [0:1]
+		res, err := client.ZRangeStore(dest, key, options.NewRangeByIndexQuery(0, 1))
+		assert.NoError(t, err)
+		res1, err := client.ZRange(dest, options.NewRangeByIndexQuery(0, 1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"a", "b"}, res1)
+		// index [0:-1] (all)
+		res, err = client.ZRangeStore(dest, key, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(3), res)
+		assert.Equal(t, []string{"a", "b", "c"}, res1)
+		// index [3:1] (none)
+		res, err = client.ZRangeStore(dest, key, options.NewRangeByIndexQuery(3, 1))
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(3, 1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(0), res)
+		assert.Equal(t, 0, len(res1))
+		// score [-inf:3]
+		var query options.ZRangeQuery
+		query = options.NewRangeByScoreQuery(
+			options.NewInfiniteScoreBoundary(options.NegativeInfinity),
+			options.NewScoreBoundary(3, true))
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, query)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(3), res)
+		assert.Equal(t, []string{"a", "b", "c"}, res1)
+		// score [-inf:3)
+		query = options.NewRangeByScoreQuery(
+			options.NewInfiniteScoreBoundary(options.NegativeInfinity),
+			options.NewScoreBoundary(3, false))
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, query)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"a", "b"}, res1)
+		// score (3:-inf] reverse
+		query = options.NewRangeByScoreQuery(
+			options.NewScoreBoundary(3, false),
+			options.NewInfiniteScoreBoundary(options.NegativeInfinity)).
+			SetReverse()
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, query)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"b", "a"}, res1)
+		// score [-inf:+inf] limit 1 2
+		query = options.NewRangeByScoreQuery(
+			options.NewInfiniteScoreBoundary(options.NegativeInfinity),
+			options.NewInfiniteScoreBoundary(options.PositiveInfinity)).
+			SetLimit(1, 2)
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"b", "c"}, res1)
+		// score [-inf:3) reverse (none)
+		query = options.NewRangeByScoreQuery(
+			options.NewInfiniteScoreBoundary(options.NegativeInfinity),
+			options.NewScoreBoundary(3, true)).
+			SetReverse()
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(0), res)
+		assert.Equal(t, 0, len(res1))
+		// score [+inf:3) (none)
+		query = options.NewRangeByScoreQuery(
+			options.NewInfiniteScoreBoundary(options.PositiveInfinity),
+			options.NewScoreBoundary(3, false))
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(0), res)
+		assert.Equal(t, 0, len(res1))
+		// lex [-:c)
+		query = options.NewRangeByLexQuery(
+			options.NewInfiniteLexBoundary(options.NegativeInfinity),
+			options.NewLexBoundary("c", false))
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"a", "b"}, res1)
+		// lex [+:-] reverse limit 1 2
+		query = options.NewRangeByLexQuery(
+			options.NewInfiniteLexBoundary(options.PositiveInfinity),
+			options.NewInfiniteLexBoundary(options.NegativeInfinity)).
+			SetReverse().SetLimit(1, 2)
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"a", "b"}, res1)
+		// lex (c:-] reverse
+		query = options.NewRangeByLexQuery(
+			options.NewLexBoundary("c", false),
+			options.NewInfiniteLexBoundary(options.NegativeInfinity)).
+			SetReverse()
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(2), res)
+		assert.Equal(t, []string{"a", "b"}, res1)
+		// lex [+:c] (none)
+		query = options.NewRangeByLexQuery(
+			options.NewInfiniteLexBoundary(options.PositiveInfinity),
+			options.NewLexBoundary("c", true))
+		res, err = client.ZRangeStore(dest, key, query)
+		assert.NoError(t, err)
+		res1, err = client.ZRange(dest, options.NewRangeByIndexQuery(0, -1))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(0), res)
+		assert.Equal(t, 0, len(res1))
+	})
+}
+
 func (suite *GlideTestSuite) TestPersist() {
 	suite.runWithDefaultClients(func(client api.BaseClient) {
 		// Test 1: Check if persist command removes the expiration time of a key.