From ddb5af9dfe829e9f557543ff70ade04791057c83 Mon Sep 17 00:00:00 2001 From: Chloe Yip <168601573+cyip10@users.noreply.github.com> Date: Mon, 19 Aug 2024 11:46:30 -0700 Subject: [PATCH] Node: add ZUNIONSTORE (#2145) * implement zunionstore Signed-off-by: Chloe Yip <chloe.yip@improving.com> * add changelog Signed-off-by: Chloe Yip <chloe.yip@improving.com> * implement zunionstore Signed-off-by: Chloe Yip <chloe.yip@improving.com> * add changelog Signed-off-by: Chloe Yip <chloe.yip@improving.com> * address comments Signed-off-by: Chloe Yip <chloe.yip@improving.com> * delete glide-for-redis submodule package Signed-off-by: Chloe Yip <chloe.yip@improving.com> * add cluster example Signed-off-by: Chloe Yip <chloe.yip@improving.com> * add remarks to base client Signed-off-by: Chloe Yip <chloe.yip@improving.com> * fix test Signed-off-by: Chloe Yip <chloe.yip@improving.com> * Apply suggestions from code review Co-authored-by: Yury-Fridlyand <yury.fridlyand@improving.com> Signed-off-by: jonathanl-bq <72158117+jonathanl-bq@users.noreply.github.com> --------- Signed-off-by: Chloe Yip <chloe.yip@improving.com> Signed-off-by: Chloe Yip <168601573+cyip10@users.noreply.github.com> Signed-off-by: jonathanl-bq <72158117+jonathanl-bq@users.noreply.github.com> Co-authored-by: jonathanl-bq <72158117+jonathanl-bq@users.noreply.github.com> Co-authored-by: Yury-Fridlyand <yury.fridlyand@improving.com> --- CHANGELOG.md | 1 + node/src/BaseClient.ts | 37 +++++ node/src/Commands.ts | 12 ++ node/src/Transaction.ts | 25 ++++ node/tests/GlideClusterClient.test.ts | 1 + node/tests/SharedTests.ts | 190 ++++++++++++++++++++++++++ node/tests/TestUtilities.ts | 2 + 7 files changed, 268 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bec770dc5..363dbbd4ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,5 @@ #### Changes +* Node: Added ZUNIONSTORE command ([#2145](https://github.com/valkey-io/valkey-glide/pull/2145)) * Node: Added XREADGROUP command ([#2124](https://github.com/valkey-io/valkey-glide/pull/2124)) * Node: Added XINFO GROUPS command ([#2122](https://github.com/valkey-io/valkey-glide/pull/2122)) * Java: Added PUBSUB CHANNELS, NUMPAT and NUMSUB commands ([#2105](https://github.com/valkey-io/valkey-glide/pull/2105)) diff --git a/node/src/BaseClient.ts b/node/src/BaseClient.ts index 3ec3081490..8160c81f8d 100644 --- a/node/src/BaseClient.ts +++ b/node/src/BaseClient.ts @@ -216,6 +216,7 @@ import { createZRevRankWithScore, createZScan, createZScore, + createZUnionStore, } from "./Commands"; import { ClosingError, @@ -3579,6 +3580,42 @@ export class BaseClient { return this.createWritePromise(createZScore(key, member)); } + /** + * Computes the union of sorted sets given by the specified `keys` and stores the result in `destination`. + * If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created. + * To get the result directly, see {@link zunionWithScores}. + * + * @see {@link https://valkey.io/commands/zunionstore/|valkey.io} for details. + * @remarks When in cluster mode, `destination` and all keys in `keys` both must map to the same hash slot. + * @param destination - The key of the destination sorted set. + * @param keys - The keys of the sorted sets with possible formats: + * string[] - for keys only. + * KeyWeight[] - for weighted keys with score multipliers. + * @param aggregationType - Specifies the aggregation strategy to apply when combining the scores of elements. See {@link AggregationType}. + * @returns The number of elements in the resulting sorted set stored at `destination`. + * + * * @example + * ```typescript + * // Example usage of zunionstore command with an existing key + * await client.zadd("key1", {"member1": 10.5, "member2": 8.2}) + * await client.zadd("key2", {"member1": 9.5}) + * await client.zunionstore("my_sorted_set", ["key1", "key2"]) // Output: 2 - Indicates that the sorted set "my_sorted_set" contains two elements. + * await client.zrangeWithScores("my_sorted_set", RangeByIndex(0, -1)) // Output: {'member1': 20, 'member2': 8.2} - "member1" is now stored in "my_sorted_set" with score of 20 and "member2" with score of 8.2. + * await client.zunionstore("my_sorted_set", ["key1", "key2"] , AggregationType.MAX ) // Output: 2 - Indicates that the sorted set "my_sorted_set" contains two elements, and each score is the maximum score between the sets. + * await client.zrangeWithScores("my_sorted_set", RangeByIndex(0, -1)) // Output: {'member1': 10.5, 'member2': 8.2} - "member1" is now stored in "my_sorted_set" with score of 10.5 and "member2" with score of 8.2. + * await client.zunionstore("my_sorted_set", ["key1, "key2], {weights: [2, 1]}) // Output: 46 + * ``` + */ + public async zunionstore( + destination: string, + keys: string[] | KeyWeight[], + aggregationType?: AggregationType, + ): Promise<number> { + return this.createWritePromise( + createZUnionStore(destination, keys, aggregationType), + ); + } + /** * Returns the scores associated with the specified `members` in the sorted set stored at `key`. * diff --git a/node/src/Commands.ts b/node/src/Commands.ts index dbe2ee674a..4947164109 100644 --- a/node/src/Commands.ts +++ b/node/src/Commands.ts @@ -1532,6 +1532,18 @@ export function createZScore( return createCommand(RequestType.ZScore, [key, member]); } +/** + * @internal + */ +export function createZUnionStore( + destination: string, + keys: string[] | KeyWeight[], + aggregationType?: AggregationType, +): command_request.Command { + const args = createZCmdStoreArgs(destination, keys, aggregationType); + return createCommand(RequestType.ZUnionStore, args); +} + /** * @internal */ diff --git a/node/src/Transaction.ts b/node/src/Transaction.ts index 794264fc20..5f0a7b7bfb 100644 --- a/node/src/Transaction.ts +++ b/node/src/Transaction.ts @@ -249,6 +249,7 @@ import { createZRevRankWithScore, createZScan, createZScore, + createZUnionStore, } from "./Commands"; import { command_request } from "./ProtobufMessage"; @@ -1769,6 +1770,30 @@ export class BaseTransaction<T extends BaseTransaction<T>> { return this.addAndReturn(createZScore(key, member)); } + /** + * Computes the union of sorted sets given by the specified `keys` and stores the result in `destination`. + * If `destination` already exists, it is overwritten. Otherwise, a new sorted set will be created. + * To get the result directly, see {@link zunionWithScores}. + * + * @see {@link https://valkey.io/commands/zunionstore/|valkey.io} for details. + * @param destination - The key of the destination sorted set. + * @param keys - The keys of the sorted sets with possible formats: + * string[] - for keys only. + * KeyWeight[] - for weighted keys with score multipliers. + * @param aggregationType - Specifies the aggregation strategy to apply when combining the scores of elements. See {@link AggregationType}. + * + * Command Response - The number of elements in the resulting sorted set stored at `destination`. + */ + public zunionstore( + destination: string, + keys: string[] | KeyWeight[], + aggregationType?: AggregationType, + ): T { + return this.addAndReturn( + createZUnionStore(destination, keys, aggregationType), + ); + } + /** * Returns the scores associated with the specified `members` in the sorted set stored at `key`. * diff --git a/node/tests/GlideClusterClient.test.ts b/node/tests/GlideClusterClient.test.ts index cadd3e5ad4..365f4b5cb6 100644 --- a/node/tests/GlideClusterClient.test.ts +++ b/node/tests/GlideClusterClient.test.ts @@ -373,6 +373,7 @@ describe("GlideClusterClient", () => { client.sinter(["abc", "zxy", "lkn"]), client.sinterstore("abc", ["zxy", "lkn"]), client.zinterstore("abc", ["zxy", "lkn"]), + client.zunionstore("abc", ["zxy", "lkn"]), client.sunionstore("abc", ["zxy", "lkn"]), client.sunion(["abc", "zxy", "lkn"]), client.pfcount(["abc", "zxy", "lkn"]), diff --git a/node/tests/SharedTests.ts b/node/tests/SharedTests.ts index 82797d3ea1..4428f6018b 100644 --- a/node/tests/SharedTests.ts +++ b/node/tests/SharedTests.ts @@ -3858,6 +3858,196 @@ export function runBaseTests<Context>(config: { config.timeout, ); + // ZUnionStore command tests + async function zunionStoreWithMaxAggregation(client: BaseClient) { + const key1 = "{testKey}:1-" + uuidv4(); + const key2 = "{testKey}:2-" + uuidv4(); + const key3 = "{testKey}:3-" + uuidv4(); + const range = { + start: 0, + stop: -1, + }; + + const membersScores1 = { one: 1.0, two: 2.0 }; + const membersScores2 = { one: 1.5, two: 2.5, three: 3.5 }; + + expect(await client.zadd(key1, membersScores1)).toEqual(2); + expect(await client.zadd(key2, membersScores2)).toEqual(3); + + // Union results are aggregated by the MAX score of elements + expect(await client.zunionstore(key3, [key1, key2], "MAX")).toEqual(3); + const zunionstoreMapMax = await client.zrangeWithScores(key3, range); + const expectedMapMax = { + one: 1.5, + two: 2.5, + three: 3.5, + }; + expect(zunionstoreMapMax).toEqual(expectedMapMax); + } + + async function zunionStoreWithMinAggregation(client: BaseClient) { + const key1 = "{testKey}:1-" + uuidv4(); + const key2 = "{testKey}:2-" + uuidv4(); + const key3 = "{testKey}:3-" + uuidv4(); + const range = { + start: 0, + stop: -1, + }; + + const membersScores1 = { one: 1.0, two: 2.0 }; + const membersScores2 = { one: 1.5, two: 2.5, three: 3.5 }; + + expect(await client.zadd(key1, membersScores1)).toEqual(2); + expect(await client.zadd(key2, membersScores2)).toEqual(3); + + // Union results are aggregated by the MIN score of elements + expect(await client.zunionstore(key3, [key1, key2], "MIN")).toEqual(3); + const zunionstoreMapMin = await client.zrangeWithScores(key3, range); + const expectedMapMin = { + one: 1.0, + two: 2.0, + three: 3.5, + }; + expect(zunionstoreMapMin).toEqual(expectedMapMin); + } + + async function zunionStoreWithSumAggregation(client: BaseClient) { + const key1 = "{testKey}:1-" + uuidv4(); + const key2 = "{testKey}:2-" + uuidv4(); + const key3 = "{testKey}:3-" + uuidv4(); + const range = { + start: 0, + stop: -1, + }; + + const membersScores1 = { one: 1.0, two: 2.0 }; + const membersScores2 = { one: 1.5, two: 2.5, three: 3.5 }; + + expect(await client.zadd(key1, membersScores1)).toEqual(2); + expect(await client.zadd(key2, membersScores2)).toEqual(3); + + // Union results are aggregated by the SUM score of elements + expect(await client.zunionstore(key3, [key1, key2], "SUM")).toEqual(3); + const zunionstoreMapSum = await client.zrangeWithScores(key3, range); + const expectedMapSum = { + one: 2.5, + two: 4.5, + three: 3.5, + }; + expect(zunionstoreMapSum).toEqual(expectedMapSum); + } + + async function zunionStoreBasicTest(client: BaseClient) { + const key1 = "{testKey}:1-" + uuidv4(); + const key2 = "{testKey}:2-" + uuidv4(); + const key3 = "{testKey}:3-" + uuidv4(); + const range = { + start: 0, + stop: -1, + }; + + const membersScores1 = { one: 1.0, two: 2.0 }; + const membersScores2 = { one: 2.0, two: 3.0, three: 4.0 }; + + expect(await client.zadd(key1, membersScores1)).toEqual(2); + expect(await client.zadd(key2, membersScores2)).toEqual(3); + + expect(await client.zunionstore(key3, [key1, key2])).toEqual(3); + const zunionstoreMap = await client.zrangeWithScores(key3, range); + const expectedMap = { + one: 3.0, + three: 4.0, + two: 5.0, + }; + expect(zunionstoreMap).toEqual(expectedMap); + } + + async function zunionStoreWithWeightsAndAggregation(client: BaseClient) { + const key1 = "{testKey}:1-" + uuidv4(); + const key2 = "{testKey}:2-" + uuidv4(); + const key3 = "{testKey}:3-" + uuidv4(); + const range = { + start: 0, + stop: -1, + }; + const membersScores1 = { one: 1.0, two: 2.0 }; + const membersScores2 = { one: 1.5, two: 2.5, three: 3.5 }; + + expect(await client.zadd(key1, membersScores1)).toEqual(2); + expect(await client.zadd(key2, membersScores2)).toEqual(3); + + // Scores are multiplied by 2.0 for key1 and key2 during aggregation. + expect( + await client.zunionstore( + key3, + [ + [key1, 2.0], + [key2, 2.0], + ], + "SUM", + ), + ).toEqual(3); + const zunionstoreMapMultiplied = await client.zrangeWithScores( + key3, + range, + ); + const expectedMapMultiplied = { + one: 5.0, + three: 7.0, + two: 9.0, + }; + expect(zunionstoreMapMultiplied).toEqual(expectedMapMultiplied); + } + + async function zunionStoreEmptyCases(client: BaseClient) { + const key1 = "{testKey}:1-" + uuidv4(); + const key2 = "{testKey}:2-" + uuidv4(); + const range = { + start: 0, + stop: -1, + }; + const membersScores1 = { one: 1.0, two: 2.0 }; + + expect(await client.zadd(key1, membersScores1)).toEqual(2); + + // Non existing key + expect( + await client.zunionstore(key2, [ + key1, + "{testKey}-non_existing_key", + ]), + ).toEqual(2); + + const zunionstore_map_nonexistingkey = await client.zrangeWithScores( + key2, + range, + ); + + const expectedMapMultiplied = { + one: 1.0, + two: 2.0, + }; + expect(zunionstore_map_nonexistingkey).toEqual(expectedMapMultiplied); + + // Empty list check + await expect(client.zunionstore("{xyz}", [])).rejects.toThrow(); + } + + it.each([ProtocolVersion.RESP2, ProtocolVersion.RESP3])( + `zunionstore test_%p`, + async (protocol) => { + await runTest(async (client: BaseClient) => { + await zunionStoreBasicTest(client); + await zunionStoreWithMaxAggregation(client); + await zunionStoreWithMinAggregation(client); + await zunionStoreWithSumAggregation(client); + await zunionStoreWithWeightsAndAggregation(client); + await zunionStoreEmptyCases(client); + }, protocol); + }, + config.timeout, + ); + it.each([ProtocolVersion.RESP2, ProtocolVersion.RESP3])( `zmscore test_%p`, async (protocol) => { diff --git a/node/tests/TestUtilities.ts b/node/tests/TestUtilities.ts index b5ae554dc7..a987b285a5 100644 --- a/node/tests/TestUtilities.ts +++ b/node/tests/TestUtilities.ts @@ -988,6 +988,8 @@ export async function transactionTest( responseData.push(["zdiffWithScores([key13, key12])", { three: 3.5 }]); baseTransaction.zdiffstore(key13, [key13, key13]); responseData.push(["zdiffstore(key13, [key13, key13])", 0]); + baseTransaction.zunionstore(key5, [key12, key13]); + responseData.push(["zunionstore(key5, [key12, key13])", 2]); baseTransaction.zmscore(key12, ["two", "one"]); responseData.push(['zmscore(key12, ["two", "one"]', [2.0, 1.0]]); baseTransaction.zinterstore(key12, [key12, key13]);