From 0b0b467c8914b5463eb80d6c11c89996563517ba Mon Sep 17 00:00:00 2001 From: Yury-Fridlyand Date: Fri, 10 Jan 2025 17:48:48 -0800 Subject: [PATCH] Go: ZRANGE (#2925) * ZRANGE Signed-off-by: Yury-Fridlyand --- go/api/base_client.go | 95 +++++++++++ go/api/options/zrange_options.go | 200 ++++++++++++++++++++++ go/api/sorted_set_commands.go | 6 +- go/integTest/glide_test_suite_test.go | 2 +- go/integTest/shared_commands_test.go | 237 ++++++++++++++++++++++++++ 5 files changed, 538 insertions(+), 2 deletions(-) create mode 100644 go/api/options/zrange_options.go diff --git a/go/api/base_client.go b/go/api/base_client.go index 61f8bef4bc..e43d664e01 100644 --- a/go/api/base_client.go +++ b/go/api/base_client.go @@ -1476,6 +1476,101 @@ func (client *baseClient) BZPopMin(keys []string, timeoutSecs float64) (Result[K return handleKeyWithMemberAndScoreResponse(result) } +// Returns the specified range of elements in the sorted set stored at `key`. +// `ZRANGE` can perform different types of range queries: by index (rank), by the score, or by lexicographical order. +// +// To get the elements with their scores, see [ZRangeWithScores]. +// +// See [valkey.io] for more details. +// +// Parameters: +// +// key - The key of the 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: +// +// An array of elements within the specified range. +// If `key` does not exist, it is treated as an empty sorted set, and the command returns an empty array. +// +// Example: +// +// // Retrieve all members of a sorted set in ascending order +// result, err := client.ZRange("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() +// result, err := client.ZRange("my_sorted_set", query) +// // `result` contains members which have scores within the range of negative infinity to 3, in descending order +// +// [valkey.io]: https://valkey.io/commands/zrange/ +func (client *baseClient) ZRange(key string, rangeQuery options.ZRangeQuery) ([]Result[string], error) { + args := make([]string, 0, 10) + args = append(args, key) + args = append(args, rangeQuery.ToArgs()...) + result, err := client.executeCommand(C.ZRange, args) + if err != nil { + return nil, err + } + + return handleStringArrayResponse(result) +} + +// Returns the specified range of elements with their scores in the sorted set stored at `key`. +// `ZRANGE` can perform different types of range queries: by index (rank), by the score, or by lexicographical order. +// +// See [valkey.io] for more details. +// +// Parameters: +// +// key - The key of the 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 score, use [RangeByScore]. +// +// Return value: +// +// A map of elements and their scores within the specified range. +// If `key` does not exist, it is treated as an empty sorted set, and the command returns an empty map. +// +// Example: +// +// // Retrieve all members of a sorted set in ascending order +// result, err := client.ZRangeWithScores("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() +// result, err := client.ZRangeWithScores("my_sorted_set", query) +// // `result` contains members with scores within the range of negative infinity to 3, in descending order +// +// [valkey.io]: https://valkey.io/commands/zrange/ +func (client *baseClient) ZRangeWithScores( + key string, + rangeQuery options.ZRangeQueryWithScores, +) (map[Result[string]]Result[float64], error) { + args := make([]string, 0, 10) + args = append(args, key) + args = append(args, rangeQuery.ToArgs()...) + args = append(args, "WITHSCORES") + result, err := client.executeCommand(C.ZRange, args) + if err != nil { + return nil, err + } + + return handleStringDoubleMapResponse(result) +} + func (client *baseClient) Persist(key string) (Result[bool], error) { result, err := client.executeCommand(C.Persist, []string{key}) if err != nil { diff --git a/go/api/options/zrange_options.go b/go/api/options/zrange_options.go new file mode 100644 index 0000000000..002dc38e24 --- /dev/null +++ b/go/api/options/zrange_options.go @@ -0,0 +1,200 @@ +// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0 + +package options + +import ( + "github.com/valkey-io/valkey-glide/go/glide/utils" +) + +// Query for `ZRange` in [SortedSetCommands] +// - For range queries by index (rank), use `RangeByIndex`. +// - For range queries by lexicographical order, use `RangeByLex`. +// - For range queries by score, use `RangeByScore`. +type ZRangeQuery interface { + ToArgs() []string +} + +// Queries a range of elements from a sorted set by theirs index. +type RangeByIndex struct { + start, end int64 + reverse bool +} + +// Queries a range of elements from a sorted set by theirs score. +type RangeByScore struct { + start, end scoreBoundary + reverse bool + Limit *Limit +} + +// Queries a range of elements from a sorted set by theirs lexicographical order. +type RangeByLex struct { + start, end lexBoundary + reverse bool + Limit *Limit +} + +type ( + InfBoundary string + scoreBoundary string + lexBoundary string +) + +const ( + // The highest bound in the sorted set + PositiveInfinity InfBoundary = "+" + // The lowest bound in the sorted set + NegativeInfinity InfBoundary = "-" +) + +// Create a new inclusive score boundary. +func NewInclusiveScoreBoundary(bound float64) scoreBoundary { + return scoreBoundary(utils.FloatToString(bound)) +} + +// Create a new score boundary. +func NewScoreBoundary(bound float64, isInclusive bool) scoreBoundary { + if !isInclusive { + return scoreBoundary("(" + utils.FloatToString(bound)) + } + return scoreBoundary(utils.FloatToString(bound)) +} + +// Create a new score boundary defined by an infinity. +func NewInfiniteScoreBoundary(bound InfBoundary) scoreBoundary { + return scoreBoundary(string(bound) + "inf") +} + +// Create a new lex boundary. +func NewLexBoundary(bound string, isInclusive bool) lexBoundary { + if !isInclusive { + return lexBoundary("(" + bound) + } + return lexBoundary("[" + bound) +} + +// Create a new lex boundary defined by an infinity. +func NewInfiniteLexBoundary(bound InfBoundary) lexBoundary { + return lexBoundary(string(bound)) +} + +// TODO re-use limit from `SORT` https://github.com/valkey-io/valkey-glide/pull/2888 +// Limit struct represents the range of elements to retrieve +// The LIMIT argument is commonly used to specify a subset of results from the matching elements, similar to the +// LIMIT clause in SQL (e.g., `SELECT LIMIT offset, count`). +type Limit struct { + // The starting position of the range, zero based. + offset int64 + // The maximum number of elements to include in the range. A negative count returns all elementsnfrom the offset. + count int64 +} + +func (limit *Limit) toArgs() []string { + return []string{"LIMIT", utils.IntToString(limit.offset), utils.IntToString(limit.count)} +} + +// Queries a range of elements from a sorted set by theirs index. +// +// Parameters: +// +// start - The start index of the range. +// end - The end index of the range. +func NewRangeByIndexQuery(start int64, end int64) *RangeByIndex { + return &RangeByIndex{start, end, false} +} + +// Reverses the sorted set, with index `0` as the element with the highest score. +func (rbi *RangeByIndex) SetReverse() *RangeByIndex { + rbi.reverse = true + return rbi +} + +func (rbi *RangeByIndex) ToArgs() []string { + args := make([]string, 0, 3) + args = append(args, utils.IntToString(rbi.start), utils.IntToString(rbi.end)) + if rbi.reverse { + args = append(args, "REV") + } + return args +} + +// Queries a range of elements from a sorted set by theirs score. +// +// Parameters: +// +// start - The start score of the range. +// end - The end score of the range. +func NewRangeByScoreQuery(start scoreBoundary, end scoreBoundary) *RangeByScore { + return &RangeByScore{start, end, false, nil} +} + +// Reverses the sorted set, with index `0` as the element with the highest score. +func (rbs *RangeByScore) SetReverse() *RangeByScore { + rbs.reverse = true + return rbs +} + +// The limit argument for a range query, unset by default. See [Limit] for more information. +func (rbs *RangeByScore) SetLimit(offset, count int64) *RangeByScore { + rbs.Limit = &Limit{offset, count} + return rbs +} + +func (rbs *RangeByScore) ToArgs() []string { + args := make([]string, 0, 7) + args = append(args, string(rbs.start), string(rbs.end), "BYSCORE") + if rbs.reverse { + args = append(args, "REV") + } + if rbs.Limit != nil { + args = append(args, rbs.Limit.toArgs()...) + } + return args +} + +// Queries a range of elements from a sorted set by theirs lexicographical order. +// +// Parameters: +// +// start - The start lex of the range. +// end - The end lex of the range. +func NewRangeByLexQuery(start lexBoundary, end lexBoundary) *RangeByLex { + return &RangeByLex{start, end, false, nil} +} + +// Reverses the sorted set, with index `0` as the element with the highest score. +func (rbl *RangeByLex) SetReverse() *RangeByLex { + rbl.reverse = true + return rbl +} + +// The limit argument for a range query, unset by default. See [Limit] for more information. +func (rbl *RangeByLex) SetLimit(offset, count int64) *RangeByLex { + rbl.Limit = &Limit{offset, count} + return rbl +} + +func (rbl *RangeByLex) ToArgs() []string { + args := make([]string, 0, 7) + args = append(args, string(rbl.start), string(rbl.end), "BYLEX") + if rbl.reverse { + args = append(args, "REV") + } + if rbl.Limit != nil { + args = append(args, rbl.Limit.toArgs()...) + } + return args +} + +// Query for `ZRangeWithScores` in [SortedSetCommands] +// - For range queries by index (rank), use `RangeByIndex`. +// - For range queries by score, use `RangeByScore`. +type ZRangeQueryWithScores interface { + // A dummy interface to distinguish queries for `ZRange` and `ZRangeWithScores` + // `ZRangeWithScores` does not support BYLEX + dummy() + ToArgs() []string +} + +func (q *RangeByIndex) dummy() {} +func (q *RangeByScore) dummy() {} diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go index 510a28a3fc..e6b18c66b8 100644 --- a/go/api/sorted_set_commands.go +++ b/go/api/sorted_set_commands.go @@ -253,7 +253,7 @@ type SortedSetCommands interface { // A `KeyWithMemberAndScore` struct containing the key where the member was popped out, the member // itself, and the member score. If no member could be popped and the `timeout` expired, returns `nil`. // - // example + // Example: // zaddResult1, err := client.ZAdd(key1, map[string]float64{"a": 1.0, "b": 1.5}) // zaddResult2, err := client.ZAdd(key2, map[string]float64{"c": 2.0}) // result, err := client.BZPopMin([]string{key1, key2}, float64(.5)) @@ -263,6 +263,10 @@ type SortedSetCommands interface { // [blocking commands]: https://github.com/valkey-io/valkey-glide/wiki/General-Concepts#blocking-commands BZPopMin(keys []string, timeoutSecs float64) (Result[KeyWithMemberAndScore], error) + ZRange(key string, rangeQuery options.ZRangeQuery) ([]Result[string], error) + + ZRangeWithScores(key string, rangeQuery options.ZRangeQueryWithScores) (map[Result[string]]Result[float64], error) + // Returns the rank of `member` in the sorted set stored at `key`, with // scores ordered from low to high, starting from `0`. // To get the rank of `member` with its score, see [ZRankWithScore]. diff --git a/go/integTest/glide_test_suite_test.go b/go/integTest/glide_test_suite_test.go index 46752041ce..fc6a5c8ff7 100644 --- a/go/integTest/glide_test_suite_test.go +++ b/go/integTest/glide_test_suite_test.go @@ -115,7 +115,7 @@ func extractAddresses(suite *GlideTestSuite, output string) []api.NodeAddress { func runClusterManager(suite *GlideTestSuite, args []string, ignoreExitCode bool) string { pythonArgs := append([]string{"../../utils/cluster_manager.py"}, args...) output, err := exec.Command("python3", pythonArgs...).CombinedOutput() - if len(output) > 0 { + if len(output) > 0 && !ignoreExitCode { suite.T().Logf("cluster_manager.py output:\n====\n%s\n====\n", string(output)) } diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go index b63ee159ba..b21a81bd2f 100644 --- a/go/integTest/shared_commands_test.go +++ b/go/integTest/shared_commands_test.go @@ -4457,6 +4457,243 @@ func (suite *GlideTestSuite) TestZRem() { }) } +func (suite *GlideTestSuite) TestZRange() { + suite.runWithDefaultClients(func(client api.BaseClient) { + t := suite.T() + 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.ZRange(key, options.NewRangeByIndexQuery(0, 1)) + expected := []api.Result[string]{ + api.CreateStringResult("a"), + api.CreateStringResult("b"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // index [0:-1] (all) + res, err = client.ZRange(key, options.NewRangeByIndexQuery(0, -1)) + expected = []api.Result[string]{ + api.CreateStringResult("a"), + api.CreateStringResult("b"), + api.CreateStringResult("c"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // index [3:1] (none) + res, err = client.ZRange(key, options.NewRangeByIndexQuery(3, 1)) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + // score [-inf:3] + var query options.ZRangeQuery + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewScoreBoundary(3, true)) + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("a"), + api.CreateStringResult("b"), + api.CreateStringResult("c"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score [-inf:3) + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewScoreBoundary(3, false)) + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("a"), + api.CreateStringResult("b"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score (3:-inf] reverse + query = options.NewRangeByScoreQuery( + options.NewScoreBoundary(3, false), + options.NewInfiniteScoreBoundary(options.NegativeInfinity)). + SetReverse() + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("b"), + api.CreateStringResult("a"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score [-inf:+inf] limit 1 2 + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewInfiniteScoreBoundary(options.PositiveInfinity)). + SetLimit(1, 2) + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("b"), + api.CreateStringResult("c"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score [-inf:3) reverse (none) + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewScoreBoundary(3, true)). + SetReverse() + res, err = client.ZRange(key, query) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + // score [+inf:3) (none) + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.PositiveInfinity), + options.NewScoreBoundary(3, false)) + res, err = client.ZRange(key, query) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + // lex [-:c) + query = options.NewRangeByLexQuery( + options.NewInfiniteLexBoundary(options.NegativeInfinity), + options.NewLexBoundary("c", false)) + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("a"), + api.CreateStringResult("b"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // lex [+:-] reverse limit 1 2 + query = options.NewRangeByLexQuery( + options.NewInfiniteLexBoundary(options.PositiveInfinity), + options.NewInfiniteLexBoundary(options.NegativeInfinity)). + SetReverse().SetLimit(1, 2) + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("b"), + api.CreateStringResult("a"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // lex (c:-] reverse + query = options.NewRangeByLexQuery( + options.NewLexBoundary("c", false), + options.NewInfiniteLexBoundary(options.NegativeInfinity)). + SetReverse() + res, err = client.ZRange(key, query) + expected = []api.Result[string]{ + api.CreateStringResult("b"), + api.CreateStringResult("a"), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // lex [+:c] (none) + query = options.NewRangeByLexQuery( + options.NewInfiniteLexBoundary(options.PositiveInfinity), + options.NewLexBoundary("c", true)) + res, err = client.ZRange(key, query) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + }) +} + +func (suite *GlideTestSuite) TestZRangeWithScores() { + suite.runWithDefaultClients(func(client api.BaseClient) { + t := suite.T() + 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.ZRangeWithScores(key, options.NewRangeByIndexQuery(0, 1)) + expected := map[api.Result[string]]api.Result[float64]{ + api.CreateStringResult("a"): api.CreateFloat64Result(1.0), + api.CreateStringResult("b"): api.CreateFloat64Result(2.0), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // index [0:-1] (all) + res, err = client.ZRangeWithScores(key, options.NewRangeByIndexQuery(0, -1)) + expected = map[api.Result[string]]api.Result[float64]{ + api.CreateStringResult("a"): api.CreateFloat64Result(1.0), + api.CreateStringResult("b"): api.CreateFloat64Result(2.0), + api.CreateStringResult("c"): api.CreateFloat64Result(3.0), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // index [3:1] (none) + res, err = client.ZRangeWithScores(key, options.NewRangeByIndexQuery(3, 1)) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + // score [-inf:3] + query := options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewScoreBoundary(3, true)) + res, err = client.ZRangeWithScores(key, query) + expected = map[api.Result[string]]api.Result[float64]{ + api.CreateStringResult("a"): api.CreateFloat64Result(1.0), + api.CreateStringResult("b"): api.CreateFloat64Result(2.0), + api.CreateStringResult("c"): api.CreateFloat64Result(3.0), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score [-inf:3) + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewScoreBoundary(3, false)) + res, err = client.ZRangeWithScores(key, query) + expected = map[api.Result[string]]api.Result[float64]{ + api.CreateStringResult("a"): api.CreateFloat64Result(1.0), + api.CreateStringResult("b"): api.CreateFloat64Result(2.0), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score (3:-inf] reverse + query = options.NewRangeByScoreQuery( + options.NewScoreBoundary(3, false), + options.NewInfiniteScoreBoundary(options.NegativeInfinity)). + SetReverse() + res, err = client.ZRangeWithScores(key, query) + expected = map[api.Result[string]]api.Result[float64]{ + api.CreateStringResult("b"): api.CreateFloat64Result(2.0), + api.CreateStringResult("a"): api.CreateFloat64Result(1.0), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score [-inf:+inf] limit 1 2 + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewInfiniteScoreBoundary(options.PositiveInfinity)). + SetLimit(1, 2) + res, err = client.ZRangeWithScores(key, query) + expected = map[api.Result[string]]api.Result[float64]{ + api.CreateStringResult("b"): api.CreateFloat64Result(2.0), + api.CreateStringResult("c"): api.CreateFloat64Result(3.0), + } + assert.NoError(t, err) + assert.Equal(t, expected, res) + // score [-inf:3) reverse (none) + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.NegativeInfinity), + options.NewScoreBoundary(3, true)). + SetReverse() + res, err = client.ZRangeWithScores(key, query) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + // score [+inf:3) (none) + query = options.NewRangeByScoreQuery( + options.NewInfiniteScoreBoundary(options.PositiveInfinity), + options.NewScoreBoundary(3, false)) + res, err = client.ZRangeWithScores(key, query) + assert.NoError(t, err) + assert.Equal(t, 0, len(res)) + }) +} + func (suite *GlideTestSuite) TestPersist() { suite.runWithDefaultClients(func(client api.BaseClient) { // Test 1: Check if persist command removes the expiration time of a key.