[Compatibility] Added ZUNION, ZUNIONSTORE commands (#833)

* Added ZUNION, ZUNIONSTORE commands * Format fix * Fixed test failure * Fixed review commands * Build fix * Review command fixes * Buid fix
microsoft · Dec 18, 2024 · d965f8a · d965f8a
1 parent f1f9e51
commit d965f8a
Show file tree

Hide file tree

Showing 17 changed files with 1,005 additions and 5 deletions.
diff --git a/libs/resources/RespCommandsDocs.json b/libs/resources/RespCommandsDocs.json
@@ -6778,5 +6778,142 @@
         "Type": "String"
       }
     ]
+  },
+  {
+    "Command": "ZUNION",
+    "Name": "ZUNION",
+    "Summary": "Returns the union of multiple sorted sets.",
+    "Group": "SortedSet",
+    "Complexity": "O(N)\u002BO(M*log(M)) with N being the sum of the sizes of the input sorted sets, and M being the number of elements in the resulting sorted set.",
+    "Arguments": [
+      {
+        "TypeDiscriminator": "RespCommandBasicArgument",
+        "Name": "NUMKEYS",
+        "DisplayText": "numkeys",
+        "Type": "Integer"
+      },
+      {
+        "TypeDiscriminator": "RespCommandKeyArgument",
+        "Name": "KEY",
+        "DisplayText": "key",
+        "Type": "Key",
+        "ArgumentFlags": "Multiple",
+        "KeySpecIndex": 0
+      },
+      {
+        "TypeDiscriminator": "RespCommandBasicArgument",
+        "Name": "WEIGHT",
+        "DisplayText": "weight",
+        "Type": "Integer",
+        "Token": "WEIGHTS",
+        "ArgumentFlags": "Optional, Multiple"
+      },
+      {
+        "TypeDiscriminator": "RespCommandContainerArgument",
+        "Name": "AGGREGATE",
+        "Type": "OneOf",
+        "Token": "AGGREGATE",
+        "ArgumentFlags": "Optional",
+        "Arguments": [
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "SUM",
+            "DisplayText": "sum",
+            "Type": "PureToken",
+            "Token": "SUM"
+          },
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "MIN",
+            "DisplayText": "min",
+            "Type": "PureToken",
+            "Token": "MIN"
+          },
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "MAX",
+            "DisplayText": "max",
+            "Type": "PureToken",
+            "Token": "MAX"
+          }
+        ]
+      },
+      {
+        "TypeDiscriminator": "RespCommandBasicArgument",
+        "Name": "WITHSCORES",
+        "DisplayText": "withscores",
+        "Type": "PureToken",
+        "Token": "WITHSCORES",
+        "ArgumentFlags": "Optional"
+      }
+    ]
+  },
+  {
+    "Command": "ZUNIONSTORE",
+    "Name": "ZUNIONSTORE",
+    "Summary": "Stores the union of multiple sorted sets in a key.",
+    "Group": "SortedSet",
+    "Complexity": "O(N)\u002BO(M log(M)) with N being the sum of the sizes of the input sorted sets, and M being the number of elements in the resulting sorted set.",
+    "Arguments": [
+      {
+        "TypeDiscriminator": "RespCommandKeyArgument",
+        "Name": "DESTINATION",
+        "DisplayText": "destination",
+        "Type": "Key",
+        "KeySpecIndex": 0
+      },
+      {
+        "TypeDiscriminator": "RespCommandBasicArgument",
+        "Name": "NUMKEYS",
+        "DisplayText": "numkeys",
+        "Type": "Integer"
+      },
+      {
+        "TypeDiscriminator": "RespCommandKeyArgument",
+        "Name": "KEY",
+        "DisplayText": "key",
+        "Type": "Key",
+        "ArgumentFlags": "Multiple",
+        "KeySpecIndex": 1
+      },
+      {
+        "TypeDiscriminator": "RespCommandBasicArgument",
+        "Name": "WEIGHT",
+        "DisplayText": "weight",
+        "Type": "Integer",
+        "Token": "WEIGHTS",
+        "ArgumentFlags": "Optional, Multiple"
+      },
+      {
+        "TypeDiscriminator": "RespCommandContainerArgument",
+        "Name": "AGGREGATE",
+        "Type": "OneOf",
+        "Token": "AGGREGATE",
+        "ArgumentFlags": "Optional",
+        "Arguments": [
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "SUM",
+            "DisplayText": "sum",
+            "Type": "PureToken",
+            "Token": "SUM"
+          },
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "MIN",
+            "DisplayText": "min",
+            "Type": "PureToken",
+            "Token": "MIN"
+          },
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "MAX",
+            "DisplayText": "max",
+            "Type": "PureToken",
+            "Token": "MAX"
+          }
+        ]
+      }
+    ]
   }
 ]
diff --git a/libs/resources/RespCommandsInfo.json b/libs/resources/RespCommandsInfo.json
@@ -4881,5 +4881,65 @@
         "Flags": "RO, Access"
       }
     ]
+  },
+  {
+    "Command": "ZUNION",
+    "Name": "ZUNION",
+    "Arity": -3,
+    "Flags": "MovableKeys, ReadOnly",
+    "AclCategories": "Read, SortedSet, Slow",
+    "KeySpecifications": [
+      {
+        "BeginSearch": {
+          "TypeDiscriminator": "BeginSearchIndex",
+          "Index": 1
+        },
+        "FindKeys": {
+          "TypeDiscriminator": "FindKeysKeyNum",
+          "KeyNumIdx": 0,
+          "FirstKey": 1,
+          "KeyStep": 1
+        },
+        "Flags": "RO, Access"
+      }
+    ]
+  },
+  {
+    "Command": "ZUNIONSTORE",
+    "Name": "ZUNIONSTORE",
+    "Arity": -4,
+    "Flags": "DenyOom, MovableKeys, Write",
+    "FirstKey": 1,
+    "LastKey": 1,
+    "Step": 1,
+    "AclCategories": "SortedSet, Slow, Write",
+    "KeySpecifications": [
+      {
+        "BeginSearch": {
+          "TypeDiscriminator": "BeginSearchIndex",
+          "Index": 1
+        },
+        "FindKeys": {
+          "TypeDiscriminator": "FindKeysRange",
+          "LastKey": 0,
+          "KeyStep": 1,
+          "Limit": 0
+        },
+        "Flags": "OW, Update"
+      },
+      {
+        "BeginSearch": {
+          "TypeDiscriminator": "BeginSearchIndex",
+          "Index": 2
+        },
+        "FindKeys": {
+          "TypeDiscriminator": "FindKeysKeyNum",
+          "KeyNumIdx": 0,
+          "FirstKey": 1,
+          "KeyStep": 1
+        },
+        "Flags": "RO, Access"
+      }
+    ]
   }
 ]
diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs
@@ -138,10 +138,17 @@ public GarnetStatus SortedSetRange(ArgSlice key, ArgSlice min, ArgSlice max, Sor
         public GarnetStatus SortedSetDifference(ArgSlice[] keys, out Dictionary<byte[], double> pairs)
             => storageSession.SortedSetDifference(keys, out pairs);
 
+        /// <inheritdoc />
+        public GarnetStatus SortedSetUnion(ReadOnlySpan<ArgSlice> keys, double[] weights, SortedSetAggregateType aggregateType, out Dictionary<byte[], double> pairs)
+            => storageSession.SortedSetUnion(keys, weights, aggregateType, out pairs);
+
         /// <inheritdoc />
         public GarnetStatus SortedSetDifferenceStore(ArgSlice destinationKey, ReadOnlySpan<ArgSlice> keys, out int count)
             => storageSession.SortedSetDifferenceStore(destinationKey, keys, out count);
 
+        public GarnetStatus SortedSetUnionStore(ArgSlice destinationKey, ReadOnlySpan<ArgSlice> keys, double[] weights, SortedSetAggregateType aggregateType, out int count)
+            => storageSession.SortedSetUnionStore(destinationKey, keys, weights, aggregateType, out count);
+
         /// <inheritdoc />
         public GarnetStatus SortedSetScan(ArgSlice key, long cursor, string match, int count, out ArgSlice[] items)
             => storageSession.ObjectScan(GarnetObjectType.SortedSet, key, cursor, match, count, out items, ref objectContext);

diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs
@@ -184,6 +184,16 @@ public GarnetStatus SortedSetDifference(ArgSlice[] keys, out Dictionary<byte[],
             return garnetApi.SortedSetDifference(keys, out pairs);
         }
 
+        /// <inheritdoc />
+        public GarnetStatus SortedSetUnion(ReadOnlySpan<ArgSlice> keys, double[] weights, SortedSetAggregateType aggregateType, out Dictionary<byte[], double> pairs)
+        {
+            foreach (var key in keys)
+            {
+                garnetApi.WATCH(key, StoreType.Object);
+            }
+            return garnetApi.SortedSetUnion(keys, weights, aggregateType, out pairs);
+        }
+
         /// <inheritdoc />
         public GarnetStatus GeoCommands(byte[] key, ref ObjectInput input, ref GarnetObjectStoreOutput outputFooter)
         {

diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs
@@ -532,6 +532,17 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi
         /// <returns>A <see cref="GarnetStatus"/> indicating the status of the operation.</returns>
         GarnetStatus SortedSetIntersectStore(ArgSlice destinationKey, ReadOnlySpan<ArgSlice> keys, double[] weights, SortedSetAggregateType aggregateType, out int count);
 
+        /// <summary>
+        /// Performs a union of multiple sorted sets and stores the result in the destination key.
+        /// </summary>
+        /// <param name="destinationKey">The key where the result will be stored.</param>
+        /// <param name="keys">The keys of the sorted sets to union.</param>
+        /// <param name="count">The number of elements in the resulting sorted set.</param>
+        /// <param name="weights">Optional weights to apply to each sorted set.</param>
+        /// <param name="aggregateType">The type of aggregation to perform (e.g., Sum, Min, Max).</param>
+        /// <returns>A <see cref="GarnetStatus"/> indicating the status of the operation.</returns>
+        GarnetStatus SortedSetUnionStore(ArgSlice destinationKey, ReadOnlySpan<ArgSlice> keys, double[] weights, SortedSetAggregateType aggregateType, out int count);
+
         #endregion
 
         #region Set Methods
@@ -1294,6 +1305,16 @@ public interface IGarnetReadApi
         /// <returns></returns>
         GarnetStatus SortedSetDifference(ArgSlice[] keys, out Dictionary<byte[], double> pairs);
 
+        /// <summary>
+        /// Performs a union of multiple sorted sets and stores the result in a dictionary.
+        /// </summary>
+        /// <param name="keys">A read-only span of ArgSlice representing the keys of the sorted sets to union.</param>
+        /// <param name="pairs">An output dictionary where the result of the union will be stored, with byte arrays as keys and doubles as values.</param>
+        /// <param name="weights">An optional array of doubles representing the weights to apply to each sorted set during the union.</param>
+        /// <param name="aggregateType">The type of aggregation to use when combining scores from the sorted sets. Defaults to <see cref="SortedSetAggregateType.Sum"/>.</param>
+        /// <returns>A <see cref="GarnetStatus"/> indicating the status of the operation.</returns>
+        GarnetStatus SortedSetUnion(ReadOnlySpan<ArgSlice> keys, double[] weights, SortedSetAggregateType aggregateType, out Dictionary<byte[], double> pairs);
+
         /// <summary>
         /// Iterates members of SortedSet key and their associated scores using a cursor,
         /// a match pattern and count parameters

diff --git a/libs/server/Resp/CmdStrings.cs b/libs/server/Resp/CmdStrings.cs
@@ -119,12 +119,12 @@ static partial class CmdStrings
         public static ReadOnlySpan<byte> LEFT => "LEFT"u8;
         public static ReadOnlySpan<byte> BYLEX => "BYLEX"u8;
         public static ReadOnlySpan<byte> REV => "REV"u8;
+        public static ReadOnlySpan<byte> LIMIT => "LIMIT"u8;
+        public static ReadOnlySpan<byte> MIN => "MIN"u8;
+        public static ReadOnlySpan<byte> MAX => "MAX"u8;
         public static ReadOnlySpan<byte> WEIGHTS => "WEIGHTS"u8;
         public static ReadOnlySpan<byte> AGGREGATE => "AGGREGATE"u8;
         public static ReadOnlySpan<byte> SUM => "SUM"u8;
-        public static ReadOnlySpan<byte> MIN => "MIN"u8;
-        public static ReadOnlySpan<byte> MAX => "MAX"u8;
-        public static ReadOnlySpan<byte> LIMIT => "LIMIT"u8;
 
         /// <summary>
         /// Response strings