Node: Add ZINTERSTORE command (valkey-io#1513)

* add zinterstore command node * Node: added zinterstore command * Node: added zinterstore command * Node: add ZINTERSTORE command * split test to functions and fix doc * change links to valkey * fix lint errors --------- Co-authored-by: Ubuntu <ubuntu@ip-172-31-20-143.eu-west-1.compute.internal>
shohamazon · Jun 8, 2024 · 2af6a18 · 2af6a18
1 parent a72a75b
commit 2af6a18
Show file tree

Hide file tree

Showing 7 changed files with 257 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,4 +1,5 @@
 #### Changes
+* 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))
 * Python: Added OBJECT IDLETIME command ([#1474](https://github.com/aws/glide-for-redis/pull/1474))

diff --git a/node/src/BaseClient.ts b/node/src/BaseClient.ts
@@ -11,7 +11,9 @@ import {
 import * as net from "net";
 import { Buffer, BufferWriter, Reader, Writer } from "protobufjs";
 import {
+    AggregationType,
     ExpireOptions,
+    KeyWeight,
     RangeByIndex,
     RangeByLex,
     RangeByScore,
@@ -82,6 +84,7 @@ import {
     createZAdd,
     createZCard,
     createZCount,
+    createZInterstore,
     createZPopMax,
     createZPopMin,
     createZRange,
@@ -1882,6 +1885,43 @@ export class BaseClient {
         );
     }
 
+    /**
+     * Computes the intersection 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 `zinter_withscores`.
+     *
+     * When in cluster mode, `destination` and all keys in `keys` must map to the same hash slot.
+     *
+     * See https://valkey.io/commands/zinterstore/ 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 zinterstore command with an existing key
+     * await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
+     * await client.zadd("key2", {"member1": 9.5})
+     * await client.zinterstore("my_sorted_set", ["key1", "key2"]) // Output: 1 - Indicates that the sorted set "my_sorted_set" contains one element.
+     * await client.zrange_withscores("my_sorted_set", RangeByIndex(0, -1)) // Output: {'member1': 20}  - "member1"  is now stored in "my_sorted_set" with score of 20.
+     * await client.zinterstore("my_sorted_set", ["key1", "key2"] , AggregationType.MAX ) // Output: 1 - Indicates that the sorted set "my_sorted_set" contains one element, and it's score is the maximum score between the sets.
+     * await client.zrange_withscores("my_sorted_set", RangeByIndex(0, -1)) // Output: {'member1': 10.5}  - "member1"  is now stored in "my_sorted_set" with score of 10.5.
+     * ```
+     */
+    public zinterstore(
+        destination: string,
+        keys: string[] | KeyWeight[],
+        aggregationType?: AggregationType,
+    ): Promise<number> {
+        return this.createWritePromise(
+            createZInterstore(destination, keys, aggregationType),
+        );
+    }
+
     /** Returns the length of the string value stored at `key`.
      * See https://redis.io/commands/strlen/ for more details.
      *

diff --git a/node/src/Commands.ts b/node/src/Commands.ts
@@ -827,6 +827,50 @@ export function createZAdd(
     return createCommand(RequestType.ZAdd, args);
 }
 
+/**
+ * `KeyWeight` - pair of variables represents a weighted key for the `ZINTERSTORE` and `ZUNIONSTORE` sorted sets commands.
+ */
+export type KeyWeight = [string, number];
+/**
+ * `AggregationType` - representing aggregation types for `ZINTERSTORE` and `ZUNIONSTORE` sorted set commands.
+ */
+export type AggregationType = "SUM" | "MIN" | "MAX";
+
+/**
+ * @internal
+ */
+export function createZInterstore(
+    destination: string,
+    keys: string[] | KeyWeight[],
+    aggregationType?: AggregationType,
+): redis_request.Command {
+    const args = createZCmdStoreArgs(destination, keys, aggregationType);
+    return createCommand(RequestType.ZInterStore, args);
+}
+
+function createZCmdStoreArgs(
+    destination: string,
+    keys: string[] | KeyWeight[],
+    aggregationType?: AggregationType,
+): string[] {
+    const args: string[] = [destination, keys.length.toString()];
+
+    if (typeof keys[0] === "string") {
+        args.push(...(keys as string[]));
+    } else {
+        const weightsKeys = keys.map(([key]) => key);
+        args.push(...(weightsKeys as string[]));
+        const weights = keys.map(([, weight]) => weight.toString());
+        args.push("WEIGHTS", ...weights);
+    }
+
+    if (aggregationType) {
+        args.push("AGGREGATE", aggregationType);
+    }
+
+    return args;
+}
+
 /**
  * @internal
  */

diff --git a/node/src/Transaction.ts b/node/src/Transaction.ts
@@ -3,8 +3,10 @@
  */
 
 import {
+    AggregationType,
     ExpireOptions,
     InfoOptions,
+    KeyWeight,
     RangeByIndex,
     RangeByLex,
     RangeByScore,
@@ -87,6 +89,7 @@ import {
     createZAdd,
     createZCard,
     createZCount,
+    createZInterstore,
     createZPopMax,
     createZPopMin,
     createZRange,
@@ -1036,6 +1039,31 @@ export class BaseTransaction<T extends BaseTransaction<T>> {
         );
     }
 
+    /**
+     * Computes the intersection 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.
+     *
+     * When in cluster mode, `destination` and all keys in `keys` must map to the same hash slot.
+     *
+     * See https://valkey.io/commands/zinterstore/ 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 zinterstore(
+        destination: string,
+        keys: string[] | KeyWeight[],
+        aggregationType?: AggregationType,
+    ): T {
+        return this.addAndReturn(
+            createZInterstore(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.
      *

diff --git a/node/tests/RedisClusterClient.test.ts b/node/tests/RedisClusterClient.test.ts
@@ -288,6 +288,7 @@ describe("RedisClusterClient", () => {
                 client.smove("abc", "zxy", "value"),
                 client.renamenx("abc", "zxy"),
                 client.sinter(["abc", "zxy", "lkn"]),
+                client.zinterstore("abc", ["zxy", "lkn"]),
                 // TODO all rest multi-key commands except ones tested below
             ];
 

diff --git a/node/tests/SharedTests.ts b/node/tests/SharedTests.ts
@@ -1878,6 +1878,141 @@ export function runBaseTests<Context>(config: {
         config.timeout,
     );
 
+    // Zinterstore command tests
+    async function zinterstoreWithAggregation(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);
+
+        // Intersection results are aggregated by the MAX score of elements
+        expect(await client.zinterstore(key3, [key1, key2], "MAX")).toEqual(2);
+        const zinterstoreMapMax = await client.zrangeWithScores(key3, range);
+        const expectedMapMax = {
+            one: 2,
+            two: 3,
+        };
+        expect(compareMaps(zinterstoreMapMax, expectedMapMax)).toBe(true);
+
+        // Intersection results are aggregated by the MIN score of elements
+        expect(await client.zinterstore(key3, [key1, key2], "MIN")).toEqual(2);
+        const zinterstoreMapMin = await client.zrangeWithScores(key3, range);
+        const expectedMapMin = {
+            one: 1,
+            two: 2,
+        };
+        expect(compareMaps(zinterstoreMapMin, expectedMapMin)).toBe(true);
+
+        // Intersection results are aggregated by the SUM score of elements
+        expect(await client.zinterstore(key3, [key1, key2], "SUM")).toEqual(2);
+        const zinterstoreMapSum = await client.zrangeWithScores(key3, range);
+        const expectedMapSum = {
+            one: 3,
+            two: 5,
+        };
+        expect(compareMaps(zinterstoreMapSum, expectedMapSum)).toBe(true);
+    }
+
+    async function zinterstoreBasicTest(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.zinterstore(key3, [key1, key2])).toEqual(2);
+        const zinterstoreMap = await client.zrangeWithScores(key3, range);
+        const expectedMap = {
+            one: 3,
+            two: 5,
+        };
+        expect(compareMaps(zinterstoreMap, expectedMap)).toBe(true);
+    }
+
+    async function zinterstoreWithWeightsAndAggregation(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);
+
+        // Scores are multiplied by 2.0 for key1 and key2 during aggregation.
+        expect(
+            await client.zinterstore(
+                key3,
+                [
+                    [key1, 2.0],
+                    [key2, 2.0],
+                ],
+                "SUM",
+            ),
+        ).toEqual(2);
+        const zinterstoreMapMultiplied = await client.zrangeWithScores(
+            key3,
+            range,
+        );
+        const expectedMapMultiplied = {
+            one: 6,
+            two: 10,
+        };
+        expect(
+            compareMaps(zinterstoreMapMultiplied, expectedMapMultiplied),
+        ).toBe(true);
+    }
+
+    async function zinterstoreEmptyCases(client: BaseClient) {
+        const key1 = "{testKey}:1-" + uuidv4();
+        const key2 = "{testKey}:2-" + uuidv4();
+
+        // Non existing key
+        expect(
+            await client.zinterstore(key2, [
+                key1,
+                "{testKey}-non_existing_key",
+            ]),
+        ).toEqual(0);
+
+        // Empty list check
+        await expect(client.zinterstore("{xyz}", [])).rejects.toThrow();
+    }
+
+    it.each([ProtocolVersion.RESP2, ProtocolVersion.RESP3])(
+        `zinterstore test_%p`,
+        async (protocol) => {
+            await runTest(async (client: BaseClient) => {
+                await zinterstoreBasicTest(client);
+                await zinterstoreWithAggregation(client);
+                await zinterstoreWithWeightsAndAggregation(client);
+                await zinterstoreEmptyCases(client);
+            }, protocol);
+        },
+        config.timeout,
+    );
+
     it.each([ProtocolVersion.RESP2, ProtocolVersion.RESP3])(
         `type test_%p`,
         async (protocol) => {

diff --git a/node/tests/TestUtilities.ts b/node/tests/TestUtilities.ts
@@ -229,6 +229,8 @@ export async function transactionTest(
     const key9 = "{key}" + uuidv4();
     const key10 = "{key}" + uuidv4();
     const key11 = "{key}" + uuidv4(); // hyper log log
+    const key12 = "{key}" + uuidv4();
+    const key13 = "{key}" + uuidv4();
     const field = uuidv4();
     const value = uuidv4();
     const args: ReturnType[] = [];
@@ -349,6 +351,12 @@ export async function transactionTest(
     args.push(["member2", "member3", "member4", "member5"]);
     baseTransaction.zrangeWithScores(key8, { start: 0, stop: -1 });
     args.push({ member2: 3, member3: 3.5, member4: 4, member5: 5 });
+    baseTransaction.zadd(key12, { one: 1, two: 2 });
+    args.push(2);
+    baseTransaction.zadd(key13, { one: 1, two: 2, tree: 3.5 });
+    args.push(3);
+    baseTransaction.zinterstore(key12, [key12, key13]);
+    args.push(2);
     baseTransaction.zcount(key8, { value: 2 }, "positiveInfinity");
     args.push(4);
     baseTransaction.zpopmin(key8);