From ad14321d274aca480420fb78604cdde9f5f6d9ff Mon Sep 17 00:00:00 2001 From: bhima2001 <56587387+bhima2001@users.noreply.github.com> Date: Fri, 25 Oct 2024 13:04:41 +0530 Subject: [PATCH] #1117 ADDED Support for ZPOPMAX (#1158) --- docs/src/content/docs/commands/ZPOPMAX.md | 148 ++++++++++++++++++ .../commands/http/zpopmax_test.go | 104 ++++++++++++ .../commands/resp/zpopmax_test.go | 75 +++++++++ .../commands/websocket/zpopmax_test.go | 75 +++++++++ internal/eval/commands.go | 13 ++ internal/eval/eval_test.go | 123 +++++++++++++++ internal/eval/sortedset/sorted_set.go | 26 +++ internal/eval/store_eval.go | 61 +++++++- internal/server/cmd_meta.go | 7 +- internal/worker/cmd_meta.go | 7 +- 10 files changed, 634 insertions(+), 5 deletions(-) create mode 100644 docs/src/content/docs/commands/ZPOPMAX.md create mode 100644 integration_tests/commands/http/zpopmax_test.go create mode 100644 integration_tests/commands/resp/zpopmax_test.go create mode 100644 integration_tests/commands/websocket/zpopmax_test.go diff --git a/docs/src/content/docs/commands/ZPOPMAX.md b/docs/src/content/docs/commands/ZPOPMAX.md new file mode 100644 index 000000000..c0fd58b7d --- /dev/null +++ b/docs/src/content/docs/commands/ZPOPMAX.md @@ -0,0 +1,148 @@ +--- +title: ZPOPMAX +description: The `ZPOPMAX` command in DiceDB is used to remove and return the members with the highest scores from the sorted set data structure at the specified key. The second argument count is optional which specifies the number of elements that needs to be popped from the sorted set. +--- + +The `ZPOPMAX` command in DiceDB is used to remove and return the members with the highest scores from the sorted set data structure at the specified key. The second argument count is optional which specifies the number of elements that needs to be popped from the sorted set. + +## Syntax + +``` +ZPOPMAX key [count] +``` + +## Parameters + +| Parameter | Description | Type | Required | +|------------|----------------------------------------------------------------------------------------------|---------|----------| +| `key` | The name of the sorted set data structure. If it does not exist, an empty array is returned. | String | Yes | +| `count` | The count argument specifies the maximum number of members to return from highest to lowest. | Integer | No | + +## Return values + +| Condition | Return Value | +|----------------------------------------------------------|------------------------------------------| +| If the key is of valid type and records are present | List of members including their scores | +| If the key does not exist or if the sorted set is empty | `(empty list or set)` | + +## Behaviour + +- The command first checks if the specified key exists. +- If the key does not exist, an empty array is returned. +- If the key exists but is not a sorted set, an error is returned. +- If the `count` argument is specified, up to that number of members with the highest scores are returned and removed. +- The returned array contains the members and their corresponding scores in the order of highest to lowest. + +## Errors +1. `Wrong type error`: + - Error Message: `(error) WRONGTYPE Operation against a key holding the wrong kind of value` + - Occurs when trying to use the command on a key that is not a sorted set. + +2. `Syntax error`: + - Error Message: `(error) ERROR wrong number of arguments for 'zpopMAX' command` + - Occurs when the command syntax is incorrect or missing required parameters. + +3. `Invalid argument type error`: + - Error Message : `(error) ERR value is not an integer or out of range` + - Occurs when the count argument passed to the command is not an integer. + +## Examples + +### Non-Existing Key (without count argument) + +Attempting to pop the member with the highest score from a non-existent sorted set: + +```bash +127.0.0.1:7379> ZPOPMAX NON_EXISTENT_KEY +(empty array) +``` + +### Existing Key (without count argument) + +Popping the member with the highest score from an existing sorted set: + +```bash +127.0.0.1:7379> ZADD myzset 1 member1 2 member2 3 member3 +(integer) 3 +127.0.0.1:7379> ZPOPMAX myzset +1) 1 "member1" +``` + +### With Count Argument + +Popping multiple members with the highest scores using the count argument: + +```bash +127.0.0.1:7379> ZADD myzset 1 member1 2 member2 3 member3 +(integer) 3 +127.0.0.1:7379> ZPOPMAX myzset 2 +1) 1 "member1" +2) 2 "member2" +``` + +### Count Argument but Multiple Members Have the Same Score + +Popping members when multiple members share the same score: + +```bash +127.0.0.1:7379> ZADD myzset 1 member1 1 member2 1 member3 +(integer) 3 +127.0.0.1:7379> ZPOPMAX myzset 2 +1) 1 "member1" +2) 1 "member2" +``` + +### Negative Count Argument + +Attempting to pop members using a negative count argument: + +```bash +127.0.0.1:7379> ZADD myzset 1 member1 2 member2 3 member3 +(integer) 3 +127.0.0.1:7379> ZPOPMAX myzset -1 +(empty array) +``` + +### Floating-Point Scores + +Popping members with floating-point scores: + +```bash +127.0.0.1:7379> ZADD myzset 1.5 member1 2.7 member2 3.8 member3 +(integer) 3 +127.0.0.1:7379> ZPOPMAX myzset +1) 1.5 "member1" +``` + +### Wrong number of arguments + +Attempting to pop from a key that is not a sorted set: + +```bash +127.0.0.1:7379> SET stringkey "string_value" +OK +127.0.0.1:7379> ZPOPMAX stringkey +(error) WRONGTYPE Operation against a key holding the wrong kind of value +``` + +### Invalid Count Argument + +Using an invalid (non-integer) count argument: + +```bash +127.0.0.1:7379> ZADD myzset 1 member1 +(integer) 1 +127.0.0.1:7379> ZPOPMAX myzset INCORRECT_COUNT_ARGUMENT +(error) ERR value is not an integer or out of range +``` + +### Wrong Type of Key (without count argument) + +Attempting to pop from a key that is not a sorted set: + +```bash +127.0.0.1:7379> SET stringkey "string_value" +OK +127.0.0.1:7379> ZPOPMAX stringkey +(error) WRONGTYPE Operation against a key holding the wrong kind of value +``` \ No newline at end of file diff --git a/integration_tests/commands/http/zpopmax_test.go b/integration_tests/commands/http/zpopmax_test.go new file mode 100644 index 000000000..971f31924 --- /dev/null +++ b/integration_tests/commands/http/zpopmax_test.go @@ -0,0 +1,104 @@ +package http + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestZPOPMAX(t *testing.T) { + exec := NewHTTPCommandExecutor() + testCases := []TestCase{ + { + name: "ZPOPMAX on non-existing key with/without count argument", + commands: []HTTPCommand{ + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "invalidTest"}}, + }, + expected: []interface{}{[]interface{}{}}, + }, + { + name: "ZPOPMAX with wrong type of key with/without count argument", + commands: []HTTPCommand{ + {Command: "SET", Body: map[string]interface{}{"key": "sortedSet", "value": "testString"}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet"}}, + }, + expected: []interface{}{"OK", "WRONGTYPE Operation against a key holding the wrong kind of value"}, + }, + { + name: "ZPOPMAX on existing key (without count argument)", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "member1", "2", "member2", "3", "member3"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet"}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "2.98"}}}, + }, + expected: []interface{}{float64(3), []interface{}{"member3", "3"}, float64(2)}, + }, + { + name: "ZPOPMAX with normal count argument", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "member1", "2", "member2", "3", "member3"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet", "value": int64(2)}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"0.44", "2"}}}, + }, + expected: []interface{}{float64(3), []interface{}{"member3", "3", "member2", "2"}, float64(1)}, + }, + { + name: "ZPOPMAX with count argument but multiple members have the same score", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "member1", "1", "member2", "1", "member3"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet", "value": int64(2)}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "2"}}}, + }, + expected: []interface{}{float64(3), []interface{}{"member3", "1", "member2", "1"}, float64(1)}, + }, + { + name: "ZPOPMAX with negative count argument", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "member1", "2", "member2", "3", "member3"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet", "value": int64(-1)}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "1000"}}}, + }, + expected: []interface{}{float64(3), []interface{}{}, float64(3)}, + }, + { + name: "ZPOPMAX with invalid count argument", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "member1"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet", "value": "INCORRECT_COUNT_ARGUMENT"}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "2"}}}, + }, + expected: []interface{}{float64(1), "ERR value is out of range, must be positive", float64(1)}, + }, + { + name: "ZPOPMAX with count argument greater than length of sorted set", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "member1", "2", "member2", "3", "member3"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet", "value": int64(10)}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1", "2"}}}, + }, + expected: []interface{}{float64(3), []interface{}{"member3", "3", "member2", "2", "member1", "1"}, float64(0)}, + }, + { + name: "ZPOPMAX with floating-point scores", + commands: []HTTPCommand{ + {Command: "ZADD", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1.5", "member1", "2.7", "member2", "3.8", "member3"}}}, + {Command: "ZPOPMAX", Body: map[string]interface{}{"key": "sortedSet"}}, + {Command: "ZCOUNT", Body: map[string]interface{}{"key": "sortedSet", "values": [...]string{"1.3", "3.6"}}}, + }, + expected: []interface{}{float64(3), []interface{}{"member3", "3.8"}, float64(2)}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.commands { + result, _ := exec.FireCommand(cmd) + assert.Equal(t, tc.expected[i], result) + } + exec.FireCommand(HTTPCommand{ + Command: "DEL", + Body: map[string]interface{}{"key": "sortedSet"}, + }) + }) + } +} diff --git a/integration_tests/commands/resp/zpopmax_test.go b/integration_tests/commands/resp/zpopmax_test.go new file mode 100644 index 000000000..f89389259 --- /dev/null +++ b/integration_tests/commands/resp/zpopmax_test.go @@ -0,0 +1,75 @@ +package resp + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestZPOPMax(t *testing.T) { + conn := getLocalConnection() + defer conn.Close() + + testCases := []TestCase{ + { + name: "ZPOPMAX on non-existing key with/without count argument", + commands: []string{"ZPOPMAX NON_EXISTENT_KEY"}, + expected: []interface{}{[]interface{}{}}, + }, + { + name: "ZPOPMAX with wrong type of key with/without count argument", + commands: []string{"SET stringkey string_value", "ZPOPMAX stringkey"}, + expected: []interface{}{"OK", "WRONGTYPE Operation against a key holding the wrong kind of value"}, + }, + { + name: "ZPOPMAX on existing key (without count argument)", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{int64(3), []interface{}{"member3", "3"}, int64(2)}, + }, + { + name: "ZPOPMAX with normal count argument", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet 2", "ZCOUNT sortedSet 1 2"}, + expected: []interface{}{int64(3), []interface{}{"member3", "3", "member2", "2"}, int64(1)}, + }, + { + name: "ZPOPMAX with count argument but multiple members have the same score", + commands: []string{"ZADD sortedSet 1 member1 1 member2 1 member3", "ZPOPMAX sortedSet 2", "ZCOUNT sortedSet 1 1"}, + expected: []interface{}{int64(3), []interface{}{"member3", "1", "member2", "1"}, int64(1)}, + }, + { + name: "ZPOPMAX with negative count argument", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet -1", "ZCOUNT sortedSet 0.6 3.231"}, + expected: []interface{}{int64(3), []interface{}{}, int64(3)}, + }, + { + name: "ZPOPMAX with invalid count argument", + commands: []string{"ZADD sortedSet 1 member1", "ZPOPMAX sortedSet INCORRECT_COUNT_ARGUMENT", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{int64(1), "ERR value is out of range, must be positive", int64(1)}, + }, + { + name: "ZPOPMAX with count argument greater than length of sorted set", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet 10", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{int64(3), []interface{}{"member3", "3", "member2", "2", "member1", "1"}, int64(0)}, + }, + { + name: "ZPOPMAX on empty sorted set", + commands: []string{"ZADD sortedSet 1 member1", "ZPOPMAX sortedSet 1", "ZPOPMAX sortedSet", "ZCOUNT sortedSet 0 10000"}, + expected: []interface{}{int64(1), []interface{}{"member1", "1"}, []interface{}{}, int64(0)}, + }, + { + name: "ZPOPMAX with floating-point scores", + commands: []string{"ZADD sortedSet 1.5 member1 2.7 member2 3.8 member3", "ZPOPMAX sortedSet", "ZCOUNT sortedSet 1.499 2.711"}, + expected: []interface{}{int64(3), []interface{}{"member3", "3.8"}, int64(2)}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.commands { + result := FireCommand(conn, cmd) + assert.Equal(t, tc.expected[i], result) + } + FireCommand(conn, "DEL sortedSet") + }) + } +} diff --git a/integration_tests/commands/websocket/zpopmax_test.go b/integration_tests/commands/websocket/zpopmax_test.go new file mode 100644 index 000000000..f91557601 --- /dev/null +++ b/integration_tests/commands/websocket/zpopmax_test.go @@ -0,0 +1,75 @@ +package websocket + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestZPOPMAX(t *testing.T) { + exec := NewWebsocketCommandExecutor() + testCases := []TestCase{ + { + name: "ZPOPMAX on non-existing key with/without count argument", + commands: []string{"ZPOPMAX NON_EXISTENT_KEY"}, + expected: []interface{}{[]interface{}{}}, + }, + { + name: "ZPOPMAX with wrong type of key with/without count argument", + commands: []string{"SET stringkey string_value", "ZPOPMAX stringkey", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{"OK", "WRONGTYPE Operation against a key holding the wrong kind of value", float64(0)}, + }, + { + name: "ZPOPMAX on existing key (without count argument)", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{float64(3), []interface{}{"member3", "3"}, float64(2)}, + }, + { + name: "ZPOPMAX with normal count argument", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet 2", "ZCOUNT sortedSet 1 2"}, + expected: []interface{}{float64(3), []interface{}{"member3", "3", "member2", "2"}, float64(1)}, + }, + { + name: "ZPOPMAX with count argument but multiple members have the same score", + commands: []string{"ZADD sortedSet 1 member1 1 member2 1 member3", "ZPOPMAX sortedSet 2", "ZCOUNT sortedSet 0.8 1"}, + expected: []interface{}{float64(3), []interface{}{"member3", "1", "member2", "1"}, float64(1)}, + }, + { + name: "ZPOPMAX with negative count argument", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet -1", "ZCOUNT sortedSet -10000 10000"}, + expected: []interface{}{float64(3), []interface{}{}, float64(3)}, + }, + { + name: "ZPOPMAX with invalid count argument", + commands: []string{"ZADD sortedSet 1 member1", "ZPOPMAX sortedSet INCORRECT_COUNT_ARGUMENT", "ZCOUNT sortedSet -1 9"}, + expected: []interface{}{float64(1), "ERR value is out of range, must be positive", float64(1)}, + }, + { + name: "ZPOPMAX with count argument greater than length of sorted set", + commands: []string{"ZADD sortedSet 1 member1 2 member2 3 member3", "ZPOPMAX sortedSet 10", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{float64(3), []interface{}{"member3", "3", "member2", "2", "member1", "1"}, float64(0)}, + }, + { + name: "ZPOPMAX on empty sorted set", + commands: []string{"ZADD sortedSet 1 member1", "ZPOPMAX sortedSet 1", "ZPOPMAX sortedSet", "ZCOUNT sortedSet 1 10"}, + expected: []interface{}{float64(1), []interface{}{"member1", "1"}, []interface{}{}, float64(0)}, + }, + { + name: "ZPOPMAX with floating-point scores", + commands: []string{"ZADD sortedSet 1.5 member1 2.7 member2 3.8 member3", "ZPOPMAX sortedSet", "ZCOUNT sortedSet 1.4999 2.700001"}, + expected: []interface{}{float64(3), []interface{}{"member3", "3.8"}, float64(2)}, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + conn := exec.ConnectToServer() + + for i, cmd := range tc.commands { + result, err := exec.FireCommandAndReadResponse(conn, cmd) + assert.Nil(t, err) + assert.Equal(t, tc.expected[i], result) + } + DeleteKey(t, conn, exec, "sortedSet") + }) + } +} diff --git a/internal/eval/commands.go b/internal/eval/commands.go index 06ac32638..b65b6f83f 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -1114,6 +1114,18 @@ var ( IsMigrated: true, NewEval: evalZRANGE, } + zpopmaxCmdMeta = DiceCmdMeta{ + Name: "ZPOPMAX", + Info: `ZPOPMAX key [count] + Pops count number of elements from the sorted set from highest to lowest and returns those score and member. + If count is not provided '1' is considered by default. + The element with the highest score is removed first + if two elements have same score then the element which is lexicographically higher is popped first`, + Arity: -1, + KeySpecs: KeySpecs{BeginIndex: 1}, + IsMigrated: true, + NewEval: evalZPOPMAX, + } zpopminCmdMeta = DiceCmdMeta{ Name: "ZPOPMIN", Info: `ZPOPMIN key [count] @@ -1335,6 +1347,7 @@ func init() { DiceCmds["ZADD"] = zaddCmdMeta DiceCmds["ZCOUNT"] = zcountCmdMeta DiceCmds["ZRANGE"] = zrangeCmdMeta + DiceCmds["ZPOPMAX"] = zpopmaxCmdMeta DiceCmds["ZPOPMIN"] = zpopminCmdMeta DiceCmds["ZRANK"] = zrankCmdMeta DiceCmds["JSON.STRAPPEND"] = jsonstrappendCmdMeta diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 32fdd3eec..31eca98e6 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -109,6 +109,7 @@ func TestEval(t *testing.T) { testEvalHRANDFIELD(t, store) testEvalZADD(t, store) testEvalZRANGE(t, store) + testEvalZPOPMAX(t, store) testEvalZPOPMIN(t, store) testEvalZRANK(t, store) testEvalHVALS(t, store) @@ -6751,6 +6752,128 @@ func BenchmarkEvalJSONSTRAPPEND(b *testing.B) { } } +func testEvalZPOPMAX(t *testing.T, store *dstore.Store) { + setup := func() { + evalZADD([]string{"myzset", "1", "member1", "2", "member2", "3", "member3"}, store) + } + + tests := map[string]evalTestCase{ + "ZPOPMAX without key": { + setup: func() {}, + input: []string{"KEY_INVALID"}, + migratedOutput: EvalResponse{ + Result: []string{}, + Error: nil, + }, + }, + "ZPOPMAX on wrongtype of key": { + setup: func() { + evalSET([]string{"mystring", "shankar"}, store) + }, + input: []string{"mystring", "1"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + }, + }, + "ZPOPMAX without count argument": { + setup: setup, + input: []string{"myzset"}, + migratedOutput: EvalResponse{ + Result: []string{"3", "member3"}, + Error: nil, + }, + }, + "ZPOPMAX with count argument": { + setup: setup, + input: []string{"myzset", "2"}, + migratedOutput: EvalResponse{ + Result: []string{"3", "member3", "2", "member2"}, + Error: nil, + }, + }, + "ZPOPMAX with count more than the elements in sorted set": { + setup: setup, + input: []string{"myzset", "4"}, + migratedOutput: EvalResponse{ + Result: []string{"3", "member3", "2", "member2", "1", "member1"}, + }, + }, + "ZPOPMAX with count as zero": { + setup: setup, + input: []string{"myzsert", "0"}, + migratedOutput: EvalResponse{ + Result: []string{}, + Error: nil, + }, + }, + "ZPOPMAX on an empty sorted set": { + setup: func() { + store.Put("myzset", store.NewObj(sortedset.New(), -1, object.ObjTypeSortedSet, object.ObjEncodingBTree)) + }, + input: []string{"myzset"}, + migratedOutput: EvalResponse{ + Result: []string{}, + Error: nil, + }, + }, + } + runMigratedEvalTests(t, tests, evalZPOPMAX, store) +} + +func BenchmarkEvalZPOPMAX(b *testing.B) { + // Define benchmark cases with varying sizes of sorted sets + benchmarks := []struct { + name string + setup func(store *dstore.Store) + input []string + }{ + { + name: "ZPOPMAX on small sorted set (10 members)", + setup: func(store *dstore.Store) { + evalZADD([]string{"sortedSet", "1", "member1", "2", "member2", "3", "member3", "4", "member4", "5", "member5", "6", "member6", "7", "member7", "8", "member8", "9", "member9", "10", "member10"}, store) + }, + input: []string{"sortedSet", "3"}, + }, + { + name: "ZPOPMAX on large sorted set (10000 members)", + setup: func(store *dstore.Store) { + args := []string{"sortedSet"} + for i := 1; i <= 10000; i++ { + args = append(args, fmt.Sprintf("%d", i), fmt.Sprintf("member%d", i)) + } + evalZADD(args, store) + }, + input: []string{"sortedSet", "10"}, + }, + { + name: "ZPOPMAX with large sorted set with duplicate scores", + setup: func(store *dstore.Store) { + args := []string{"sortedSet"} + for i := 1; i <= 10000; i++ { + args = append(args, "1", fmt.Sprintf("member%d", i)) + } + evalZADD(args, store) + }, + input: []string{"sortedSet", "2"}, + }, + } + + store := dstore.NewStore(nil, nil) + + for _, bm := range benchmarks { + b.Run(bm.name, func(b *testing.B) { + bm.setup(store) + + for i := 0; i < b.N; i++ { + // Reset the store before each run to avoid contamination + dstore.ResetStore(store) + bm.setup(store) + evalZPOPMAX(bm.input, store) + } + }) + } +} func BenchmarkZCOUNT(b *testing.B) { store := dstore.NewStore(nil, nil) diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index 15888571f..36caee5de 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -203,6 +203,32 @@ func (ss *Set) Get(member string) (float64, bool) { return score, exists } +// This func is used to remove the maximum element from the sortedset. +// It takes count as an argument which tells the number of elements to be removed from the sortedset. +func (ss *Set) PopMax(count int) []string { + result := make([]string, 2*count) + + size := 0 + for i := 0; i < count; i++ { + item := ss.tree.DeleteMax() + if item == nil { + break + } + ssi := item.(*Item) + result[2*i] = ssi.Member + result[2*i+1] = strconv.FormatFloat(ssi.Score, 'g', -1, 64) + + delete(ss.memberMap, ssi.Member) + size++ + } + + if size < count { + result = result[:2*size] + } + + return result +} + // Iterate over elements in the B-Tree with scores in the [min, max] range func (ss *Set) CountInRange(minVal, maxVal float64) int { count := 0 diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 20912edc5..9bd1d8d31 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -731,7 +731,7 @@ func evalZRANK(args []string, store *dstore.Store) *EvalResponse { } if withScore { - scoreStr := strconv.FormatFloat(score, 'f', -1, 64) + scoreStr := strconv.FormatFloat(score, 'f', -1, 64) return &EvalResponse{ Result: []interface{}{rank, scoreStr}, Error: nil, @@ -1897,7 +1897,7 @@ func evalHSCAN(args []string, store *dstore.Store) *EvalResponse { Error: nil, } } - + // evalBF.RESERVE evaluates the BF.RESERVE command responsible for initializing a // new bloom filter and allocation it's relevant parameters based on given inputs. // If no params are provided, it uses defaults. @@ -1987,3 +1987,60 @@ func evalBFINFO(args []string, store *dstore.Store) *EvalResponse { return makeEvalResult(result) } + +// This command removes the element with the maximum score from the sorted set. +// If two elements have the same score then the members are aligned in lexicographically and the lexicographically greater element is removed. +// There is a second optional element called count which specifies the number of element to be removed. +// Returns the removed elements from the sorted set. +func evalZPOPMAX(args []string, store *dstore.Store) *EvalResponse { + if len(args) < 1 || len(args) > 2 { + return &EvalResponse{ + Result: clientio.NIL, + Error: diceerrors.ErrWrongArgumentCount("ZPOPMAX"), + } + } + + key := args[0] + obj := store.Get(key) + + count := 1 + if len(args) > 1 { + ops, err := strconv.Atoi(args[1]) + if err != nil { + return &EvalResponse{ + Result: clientio.NIL, + Error: diceerrors.ErrGeneral("value is out of range, must be positive"), // This error is thrown when then count argument is not an integer + } + } + if ops <= 0 { + return &EvalResponse{ + Result: []string{}, // Returns empty array when the count is less than or equal to 0 + Error: nil, + } + } + count = ops + } + + if obj == nil { + return &EvalResponse{ + Result: []string{}, // Returns empty array when the object with given key is not present in the store + Error: nil, + } + } + + var sortedSet *sortedset.Set + sortedSet, err := sortedset.FromObject(obj) + if err != nil { + return &EvalResponse{ + Result: clientio.NIL, + Error: diceerrors.ErrWrongTypeOperation, // Returns this error when a key is present in the store but is not of type sortedset.Set + } + } + + var res []string = sortedSet.PopMax(count) + + return &EvalResponse{ + Result: res, + Error: nil, + } +} diff --git a/internal/server/cmd_meta.go b/internal/server/cmd_meta.go index 6583489c3..5714632d2 100644 --- a/internal/server/cmd_meta.go +++ b/internal/server/cmd_meta.go @@ -117,7 +117,7 @@ var ( Cmd: "JSON.OBJLEN", CmdType: SingleShard, } - hlenCmdMeta = CmdsMeta{ + hlenCmdMeta = CmdsMeta{ Cmd: "HLEN", CmdType: SingleShard, } @@ -161,6 +161,10 @@ var ( Cmd: "HRANDFIELD", CmdType: SingleShard, } + zpopmaxCmdMeta = CmdsMeta{ + Cmd: "ZPOPMAX", + CmdType: SingleShard, + } bfaddCmdMeta = CmdsMeta{ Cmd: "BF.ADD", CmdType: SingleShard, @@ -232,6 +236,7 @@ func init() { WorkerCmdsMeta["HINCRBY"] = hincrbyCmdMeta WorkerCmdsMeta["HINCRBYFLOAT"] = hincrbyfloatCmdMeta WorkerCmdsMeta["HRANDFIELD"] = hrandfieldCmdMeta + WorkerCmdsMeta["ZPOPMAX"] = zpopmaxCmdMeta WorkerCmdsMeta["BF.ADD"] = bfaddCmdMeta WorkerCmdsMeta["BF.RESERVE"] = bfreserveCmdMeta WorkerCmdsMeta["BF.EXISTS"] = bfexistsCmdMeta diff --git a/internal/worker/cmd_meta.go b/internal/worker/cmd_meta.go index 56f15c4a4..85563924c 100644 --- a/internal/worker/cmd_meta.go +++ b/internal/worker/cmd_meta.go @@ -85,6 +85,7 @@ const ( CmdHRandField = "HRANDFIELD" CmdGetRange = "GETRANGE" CmdAppend = "APPEND" + CmdZPopMax = "ZPOPMAX" CmdHLen = "HLEN" CmdHStrLen = "HSTRLEN" CmdHScan = "HSCAN" @@ -92,7 +93,6 @@ const ( CmdBFReserve = "BF.RESERVE" CmdBFInfo = "BF.INFO" CmdBFExists = "BF.EXISTS" - ) type CmdMeta struct { @@ -168,7 +168,7 @@ var CommandsMeta = map[string]CmdMeta{ CmdType: SingleShard, }, CmdHScan: { - CmdType: SingleShard, + CmdType: SingleShard, }, CmdHIncrBy: { CmdType: SingleShard, @@ -259,6 +259,9 @@ var CommandsMeta = map[string]CmdMeta{ CmdZPopMin: { CmdType: SingleShard, }, + CmdZPopMax: { + CmdType: SingleShard, + }, // Bloom Filter CmdBFAdd: {