Node: Add ZUNIONSTORE command

CHANGELOG.md

@@ -1,4 +1,5 @@
 #### Changes
+* Node: Added ZUNIONSTORE command ([#1550](https://github.com/aws/glide-for-redis/pull/1550))
 * Node: Added ZINTERSTORE command ([#1513](https://github.com/aws/glide-for-redis/pull/1513))
 * Python: Added OBJECT ENCODING command ([#1471](https://github.com/aws/glide-for-redis/pull/1471))
 * Python: Added OBJECT FREQ command ([#1472](https://github.com/aws/glide-for-redis/pull/1472))

node/src/BaseClient.ts

@@ -77,17 +77,20 @@ import {
     createSMove,
     createSPop,
     createSRem,
+    createSUnionStore,
     createSet,
     createStrlen,
     createTTL,
     createType,
     createUnlink,
     createXAdd,
+    createXLen,
     createXRead,
     createXTrim,
     createZAdd,
     createZCard,
     createZCount,
+    createZInterCard,
     createZInterstore,
     createZPopMax,
     createZPopMin,
@@ -98,9 +101,7 @@ import {
     createZRemRangeByRank,
     createZRemRangeByScore,
     createZScore,
-    createSUnionStore,
-    createXLen,
-    createZInterCard,
+    createZUnionStore,
 } from "./Commands";
 import {
     ClosingError,
@@ -1973,6 +1974,43 @@ export class BaseClient {
         );
     }
 
+    /**
+     * 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 `zunion_withscores`.
+     * 
+     * When in cluster mode, `destination` and all keys in `keys` must map to the same hash slot.
+     * 
+     * See https://valkey.io/commands/zunionstore/ for more 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 `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.
+     * ```
+     */
+    public zunionStore(
+        destination: string,
+        keys: string[] | KeyWeight[],
+        aggregationType?: AggregationType,
+    ): Promise<number> {
+        return this.createWritePromise(
+            createZUnionStore(destination, keys, aggregationType),
+        );
+    }
+
     /** Returns the length of the string value stored at `key`.
      * See https://redis.io/commands/strlen/ for more details.
      *

node/src/Commands.ts

@@ -858,6 +858,18 @@ export function createZInterstore(
     return createCommand(RequestType.ZInterStore, args);
 }
 
+/**
+ * @internal
+ */
+export function createZUnionStore(
+    destination: string,
+    keys: string[] | KeyWeight[],
+    aggregationType?: AggregationType,
+): redis_request.Command {
+    const args = createZCmdStoreArgs(destination, keys, aggregationType);
+    return createCommand(RequestType.ZUnionStore, args);
+}
+
 function createZCmdStoreArgs(
     destination: string,
     keys: string[] | KeyWeight[],

node/src/Transaction.ts

@@ -80,6 +80,7 @@ import {
     createSMove,
     createSPop,
     createSRem,
+    createSUnionStore,
     createSelect,
     createSet,
     createStrlen,
@@ -88,11 +89,13 @@ import {
     createType,
     createUnlink,
     createXAdd,
+    createXLen,
     createXRead,
     createXTrim,
     createZAdd,
     createZCard,
     createZCount,
+    createZInterCard,
     createZInterstore,
     createZPopMax,
     createZPopMin,
@@ -103,9 +106,7 @@ import {
     createZRemRangeByRank,
     createZRemRangeByScore,
     createZScore,
-    createSUnionStore,
-    createXLen,
-    createZInterCard,
+    createZUnionStore,
 } from "./Commands";
 import { redis_request } from "./ProtobufMessage";
 
@@ -1103,6 +1104,30 @@ export class BaseTransaction<T extends BaseTransaction<T>> {
         );
     }
 
+    /**
+     * 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 `zunion_withscores`.
+     * 
+     * See https://valkey.io/commands/zunionstore/ for more 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 `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 string representation of the type of the value stored at `key`.
      * See https://redis.io/commands/type/ for more details.
      *

node/tests/RedisClusterClient.test.ts

@@ -292,6 +292,7 @@ describe("RedisClusterClient", () => {
                 client.renamenx("abc", "zxy"),
                 client.sinter(["abc", "zxy", "lkn"]),
                 client.zinterstore("abc", ["zxy", "lkn"]),
+                client.zunionStore("abc", ["zxy", "lkn"]),
                 client.sunionstore("abc", ["zxy", "lkn"]),
                 client.pfcount(["abc", "zxy", "lkn"]),
                 // TODO all rest multi-key commands except ones tested below

node/tests/SharedTests.ts

@@ -2117,6 +2117,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])(
         `type test_%p`,
         async (protocol) => {