From 672c5489b2c2b281ac0433ee0e0e96eb927bd396 Mon Sep 17 00:00:00 2001 From: prateek-kumar-improving Date: Mon, 3 Feb 2025 11:58:26 -0800 Subject: [PATCH] Go: Add ZDiffStore command (#3034) * Go: Add ZDiffStore command Signed-off-by: Prateek Kumar --- go/api/base_client.go | 40 +++++ go/api/sorted_set_commands.go | 2 + go/integTest/shared_commands_test.go | 216 ++++++++++++++++++--------- 3 files changed, 188 insertions(+), 70 deletions(-) diff --git a/go/api/base_client.go b/go/api/base_client.go index c10ae51c18..0b0763dde0 100644 --- a/go/api/base_client.go +++ b/go/api/base_client.go @@ -7265,3 +7265,43 @@ func (client *baseClient) ZDiffWithScores(keys []string) (map[string]float64, er } return handleStringDoubleMapResponse(result) } + +// Calculates the difference between the first sorted set and all the successive sorted sets at +// `keys` and stores the difference as a sorted set to `destination`, +// overwriting it if it already exists. Non-existent keys are treated as empty sets. +// +// Note: When in cluster mode, `destination` and all `keys` must map to the same hash slot. +// +// Available for Valkey 6.2 and above. +// +// See [valkey.io] for details. +// +// Parameters: +// +// destination - The key for the resulting sorted set. +// keys - The keys of the sorted sets to compare. +// +// Return value: +// +// The number of members in the resulting sorted set stored at `destination`. +// +// Example: +// +// membersScores1 := map[string]float64{"one": 1.0, "two": 2.0, "three": 3.0} +// membersScores2 := map[string]float64{"two": 2.0} +// zAddResult1, err := client.ZAdd("key1", membersScores1) +// zAddResult2, err := client.ZAdd("key2", membersScores2) +// zDiffStoreResult, err := client.ZDiffStore("key4", []string{"key1", "key2"}) +// fmt.Println(zDiffStoreResult) // Output: 2 +// +// [valkey.io]: https://valkey.io/commands/zdiffstore/ +func (client *baseClient) ZDiffStore(destination string, keys []string) (int64, error) { + result, err := client.executeCommand( + C.ZDiffStore, + append([]string{destination, strconv.Itoa(len(keys))}, keys...), + ) + if err != nil { + return defaultIntResponse, err + } + return handleIntResponse(result) +} diff --git a/go/api/sorted_set_commands.go b/go/api/sorted_set_commands.go index ce010c39ca..c05ab4a8da 100644 --- a/go/api/sorted_set_commands.go +++ b/go/api/sorted_set_commands.go @@ -86,4 +86,6 @@ type SortedSetCommands interface { ZInter(keys options.KeyArray) ([]string, error) ZInterWithScores(options *options.ZInterOptions) (map[string]float64, error) + + ZDiffStore(destination string, keys []string) (int64, error) } diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go index 8fd6124d38..4cd016b005 100644 --- a/go/integTest/shared_commands_test.go +++ b/go/integTest/shared_commands_test.go @@ -7807,76 +7807,6 @@ func (suite *GlideTestSuite) TestBitFieldRO_MultipleGets() { }) } -func (suite *GlideTestSuite) TestZDiff() { - suite.runWithDefaultClients(func(client api.BaseClient) { - suite.SkipIfServerVersionLowerThanBy("6.2.0") - t := suite.T() - key1 := "{testKey}:1-" + uuid.NewString() - key2 := "{testKey}:2-" + uuid.NewString() - key3 := "{testKey}:3-" + uuid.NewString() - nonExistentKey := "{testKey}:4-" + uuid.NewString() - - membersScores1 := map[string]float64{ - "one": 1.0, - "two": 2.0, - "three": 3.0, - } - - membersScores2 := map[string]float64{ - "two": 2.0, - } - - membersScores3 := map[string]float64{ - "one": 1.0, - "two": 2.0, - "three": 3.0, - "four": 4.0, - } - - zAddResult1, err := client.ZAdd(key1, membersScores1) - assert.NoError(t, err) - assert.Equal(t, int64(3), zAddResult1) - zAddResult2, err := client.ZAdd(key2, membersScores2) - assert.NoError(t, err) - assert.Equal(t, int64(1), zAddResult2) - zAddResult3, err := client.ZAdd(key3, membersScores3) - assert.NoError(t, err) - assert.Equal(t, int64(4), zAddResult3) - - zDiffResult, err := client.ZDiff([]string{key1, key2}) - assert.NoError(t, err) - assert.Equal(t, []string{"one", "three"}, zDiffResult) - zDiffResult, err = client.ZDiff([]string{key1, key3}) - assert.NoError(t, err) - assert.Equal(t, []string{}, zDiffResult) - zDiffResult, err = client.ZDiff([]string{nonExistentKey, key3}) - assert.NoError(t, err) - assert.Equal(t, []string{}, zDiffResult) - - zDiffResultWithScores, err := client.ZDiffWithScores([]string{key1, key2}) - assert.NoError(t, err) - assert.Equal(t, map[string]float64{"one": 1.0, "three": 3.0}, zDiffResultWithScores) - zDiffResultWithScores, err = client.ZDiffWithScores([]string{key1, key3}) - assert.NoError(t, err) - assert.Equal(t, map[string]float64{}, zDiffResultWithScores) - zDiffResultWithScores, err = client.ZDiffWithScores([]string{nonExistentKey, key3}) - assert.NoError(t, err) - assert.Equal(t, map[string]float64{}, zDiffResultWithScores) - - // Key exists, but it is not a set - setResult, _ := client.Set(nonExistentKey, "bar") - assert.Equal(t, setResult, "OK") - - _, err = client.ZDiff([]string{nonExistentKey, key2}) - assert.NotNil(t, err) - assert.IsType(t, &errors.RequestError{}, err) - - _, err = client.ZDiffWithScores([]string{nonExistentKey, key2}) - assert.NotNil(t, err) - assert.IsType(t, &errors.RequestError{}, err) - }) -} - func (suite *GlideTestSuite) TestZInter() { suite.SkipIfServerVersionLowerThanBy("6.2.0") suite.runWithDefaultClients(func(client api.BaseClient) { @@ -7975,3 +7905,149 @@ func (suite *GlideTestSuite) TestZInter() { assert.IsType(suite.T(), &errors.RequestError{}, err) }) } + +func (suite *GlideTestSuite) TestZDiff() { + suite.runWithDefaultClients(func(client api.BaseClient) { + suite.SkipIfServerVersionLowerThanBy("6.2.0") + t := suite.T() + key1 := "{testKey}:1-" + uuid.NewString() + key2 := "{testKey}:2-" + uuid.NewString() + key3 := "{testKey}:3-" + uuid.NewString() + nonExistentKey := "{testKey}:4-" + uuid.NewString() + + membersScores1 := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + } + + membersScores2 := map[string]float64{ + "two": 2.0, + } + + membersScores3 := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + "four": 4.0, + } + + zAddResult1, err := client.ZAdd(key1, membersScores1) + assert.NoError(t, err) + assert.Equal(t, int64(3), zAddResult1) + zAddResult2, err := client.ZAdd(key2, membersScores2) + assert.NoError(t, err) + assert.Equal(t, int64(1), zAddResult2) + zAddResult3, err := client.ZAdd(key3, membersScores3) + assert.NoError(t, err) + assert.Equal(t, int64(4), zAddResult3) + + zDiffResult, err := client.ZDiff([]string{key1, key2}) + assert.NoError(t, err) + assert.Equal(t, []string{"one", "three"}, zDiffResult) + zDiffResult, err = client.ZDiff([]string{key1, key3}) + assert.NoError(t, err) + assert.Equal(t, []string{}, zDiffResult) + zDiffResult, err = client.ZDiff([]string{nonExistentKey, key3}) + assert.NoError(t, err) + assert.Equal(t, []string{}, zDiffResult) + + zDiffResultWithScores, err := client.ZDiffWithScores([]string{key1, key2}) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{"one": 1.0, "three": 3.0}, zDiffResultWithScores) + zDiffResultWithScores, err = client.ZDiffWithScores([]string{key1, key3}) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{}, zDiffResultWithScores) + zDiffResultWithScores, err = client.ZDiffWithScores([]string{nonExistentKey, key3}) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{}, zDiffResultWithScores) + + // Key exists, but it is not a set + setResult, _ := client.Set(nonExistentKey, "bar") + assert.Equal(t, setResult, "OK") + + _, err = client.ZDiff([]string{nonExistentKey, key2}) + assert.NotNil(t, err) + assert.IsType(t, &errors.RequestError{}, err) + + _, err = client.ZDiffWithScores([]string{nonExistentKey, key2}) + assert.NotNil(t, err) + assert.IsType(t, &errors.RequestError{}, err) + }) +} + +func (suite *GlideTestSuite) TestZDiffStore() { + suite.runWithDefaultClients(func(client api.BaseClient) { + suite.SkipIfServerVersionLowerThanBy("6.2.0") + t := suite.T() + key1 := "{testKey}:1-" + uuid.NewString() + key2 := "{testKey}:2-" + uuid.NewString() + key3 := "{testKey}:3-" + uuid.NewString() + key4 := "{testKey}:4-" + uuid.NewString() + key5 := "{testKey}:5-" + uuid.NewString() + + membersScores1 := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + } + + membersScores2 := map[string]float64{ + "two": 2.0, + } + + membersScores3 := map[string]float64{ + "one": 1.0, + "two": 2.0, + "three": 3.0, + "four": 4.0, + } + + zAddResult1, err := client.ZAdd(key1, membersScores1) + assert.NoError(t, err) + assert.Equal(t, int64(3), zAddResult1) + zAddResult2, err := client.ZAdd(key2, membersScores2) + assert.NoError(t, err) + assert.Equal(t, int64(1), zAddResult2) + zAddResult3, err := client.ZAdd(key3, membersScores3) + assert.NoError(t, err) + assert.Equal(t, int64(4), zAddResult3) + + zDiffStoreResult, err := client.ZDiffStore(key4, []string{key1, key2}) + assert.NoError(t, err) + assert.Equal(t, zDiffStoreResult, int64(2)) + zRangeWithScoreResult, err := client.ZRangeWithScores(key4, options.NewRangeByIndexQuery(0, -1)) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{"one": 1.0, "three": 3.0}, zRangeWithScoreResult) + + zDiffStoreResult, err = client.ZDiffStore(key4, []string{key3, key2, key1}) + assert.NoError(t, err) + assert.Equal(t, zDiffStoreResult, int64(1)) + zRangeWithScoreResult, err = client.ZRangeWithScores(key4, options.NewRangeByIndexQuery(0, -1)) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{"four": 4.0}, zRangeWithScoreResult) + + zDiffStoreResult, err = client.ZDiffStore(key4, []string{key1, key3}) + assert.NoError(t, err) + assert.Equal(t, zDiffStoreResult, int64(0)) + zRangeWithScoreResult, err = client.ZRangeWithScores(key4, options.NewRangeByIndexQuery(0, -1)) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{}, zRangeWithScoreResult) + + // Non-Existing key + zDiffStoreResult, err = client.ZDiffStore(key4, []string{key5, key1}) + assert.NoError(t, err) + assert.Equal(t, zDiffStoreResult, int64(0)) + zRangeWithScoreResult, err = client.ZRangeWithScores(key4, options.NewRangeByIndexQuery(0, -1)) + assert.NoError(t, err) + assert.Equal(t, map[string]float64{}, zRangeWithScoreResult) + + // Key exists, but it is not a set + setResult, err := client.Set(key5, "bar") + assert.NoError(t, err) + assert.Equal(t, setResult, "OK") + _, err = client.ZDiffStore(key4, []string{key5, key1}) + assert.NotNil(t, err) + assert.IsType(t, &errors.RequestError{}, err) + }) +}