diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs index 2b5c730ebd..6123b6b234 100644 --- a/libs/server/API/GarnetApiObjectCommands.cs +++ b/libs/server/API/GarnetApiObjectCommands.cs @@ -304,6 +304,10 @@ public GarnetStatus SetRandomMember(byte[] key, ArgSlice input, ref GarnetObject public GarnetStatus SetScan(ArgSlice key, long cursor, string match, int count, out ArgSlice[] items) => storageSession.SetScan(key, cursor, match, count, out items, ref objectContext); + /// + public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output) + => storageSession.SetUnion(keys, out output, ref objectContext); + #endregion #region Hash Methods diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index bbc0703308..0d81c4e7c3 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -251,6 +251,17 @@ public GarnetStatus SetScan(ArgSlice key, long cursor, string match, int count, return garnetApi.SetScan(key, cursor, match, count, out items); } + /// + public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output) + { + foreach (var key in keys) + { + garnetApi.WATCH(key, StoreType.Object); + } + return garnetApi.SetUnion(keys, out output); + } + + #endregion #region Hash Methods diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index b10730baeb..b011833587 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -1199,6 +1199,15 @@ public interface IGarnetReadApi /// GarnetStatus SetScan(ArgSlice key, long cursor, string match, int count, out ArgSlice[] items); + /// + /// Returns the members of the set resulting from the union of all the given sets. + /// Keys that do not exist are considered to be empty sets. + /// + /// + /// + /// + GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output); + #endregion #region Hash Methods diff --git a/libs/server/Objects/Set/SetObject.cs b/libs/server/Objects/Set/SetObject.cs index b7d983aa82..c84d8332e5 100644 --- a/libs/server/Objects/Set/SetObject.cs +++ b/libs/server/Objects/Set/SetObject.cs @@ -25,6 +25,7 @@ public enum SetOperation : byte SSCAN, SRANDMEMBER, SISMEMBER, + SUNION, } @@ -196,5 +197,7 @@ public override unsafe void Scan(long start, out List items, out long cu if (cursor == set.Count) cursor = 0; } + + public HashSet Set => set; } } \ No newline at end of file diff --git a/libs/server/Resp/Objects/SetCommands.cs b/libs/server/Resp/Objects/SetCommands.cs index bb4b14ea45..636a5f542c 100644 --- a/libs/server/Resp/Objects/SetCommands.cs +++ b/libs/server/Resp/Objects/SetCommands.cs @@ -94,6 +94,58 @@ private unsafe bool SetAdd(int count, byte* ptr, ref TGarnetApi stor return true; } + /// + /// Returns the members of the set resulting from the union of all the given sets. + /// Keys that do not exist are considered to be empty sets. + /// + /// + /// + /// + /// + /// + private bool SetUnion(int count, byte* ptr, ref TGarnetApi storageApi) + where TGarnetApi : IGarnetApi + { + if (count < 1) + { + return AbortWithWrongNumberOfArguments("SUNION", count); + } + + // Read all the keys + ArgSlice[] keys = new ArgSlice[count]; + + for (int i = 0; i < keys.Length; i++) + { + keys[i] = default; + if (!RespReadUtils.ReadPtrWithLengthHeader(ref keys[i].ptr, ref keys[i].length, ref ptr, recvBufferPtr + bytesRead)) + return false; + } + + if (NetworkKeyArraySlotVerify(ref keys, true)) + { + var bufSpan = new ReadOnlySpan(recvBufferPtr, bytesRead); + if (!DrainCommands(bufSpan, count)) return false; + return true; + } + + storageApi.SetUnion(keys, out var result); + + // write the size of result + var resultCount = result.Count; + while (!RespWriteUtils.WriteArrayLength(resultCount, ref dcurr, dend)) + SendAndReset(); + + foreach (var item in result) + { + while (!RespWriteUtils.WriteBulkString(item, ref dcurr, dend)) + SendAndReset(); + } + + // update read pointers + readHead = (int)(ptr - recvBufferPtr); + return true; + } + /// /// Remove the specified members from the set. /// Specified members that are not a member of this set are ignored. diff --git a/libs/server/Resp/RespCommand.cs b/libs/server/Resp/RespCommand.cs index 672defd770..a0b9a2113c 100644 --- a/libs/server/Resp/RespCommand.cs +++ b/libs/server/Resp/RespCommand.cs @@ -693,6 +693,10 @@ static RespCommand MatchedNone(RespServerSession session, int oldReadHead) { return (RespCommand.STRLEN, 0); } + else if (*(ulong*)(ptr + 4) == MemoryMarshal.Read("SUNION\r\n"u8)) + { + return (RespCommand.Set, (byte)SetOperation.SUNION); + } break; case 'U': diff --git a/libs/server/Resp/RespCommandsInfo.cs b/libs/server/Resp/RespCommandsInfo.cs index 0700107984..97ddb14e35 100644 --- a/libs/server/Resp/RespCommandsInfo.cs +++ b/libs/server/Resp/RespCommandsInfo.cs @@ -220,6 +220,7 @@ public static RespCommandsInfo findCommand(RespCommand cmd, byte subCmd = 0) {(byte)SetOperation.SPOP, new RespCommandsInfo("SPOP", RespCommand.Set, -1, null, (byte)SetOperation.SPOP) }, {(byte)SetOperation.SSCAN, new RespCommandsInfo("SSCAN", RespCommand.Set, -2, null, (byte)SetOperation.SSCAN) }, {(byte)SetOperation.SISMEMBER, new RespCommandsInfo("SISMEMBER",RespCommand.Set, 2, null, (byte)SetOperation.SISMEMBER) }, + {(byte)SetOperation.SUNION, new RespCommandsInfo("SUNION", RespCommand.Set, -1, null, (byte)SetOperation.SUNION) }, }; private static readonly Dictionary customCommandsInfoMap = new Dictionary diff --git a/libs/server/Resp/RespInfo.cs b/libs/server/Resp/RespInfo.cs index 2b81b58c51..81847d6863 100644 --- a/libs/server/Resp/RespInfo.cs +++ b/libs/server/Resp/RespInfo.cs @@ -38,7 +38,7 @@ public static HashSet GetCommands() // Pub/sub "PUBLISH", "SUBSCRIBE", "PSUBSCRIBE", "UNSUBSCRIBE", "PUNSUBSCRIBE", // Set - "SADD", "SREM", "SPOP", "SMEMBERS", "SCARD", "SSCAN", "SRANDMEMBER", "SISMEMBER", + "SADD", "SREM", "SPOP", "SMEMBERS", "SCARD", "SSCAN", "SRANDMEMBER", "SISMEMBER", "SUNION", //Scan ops "DBSIZE", "KEYS","SCAN", // Geospatial commands diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index 6dea66abdc..07b934a792 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -508,6 +508,7 @@ private bool ProcessArrayCommands(RespCommand cmd, byte subcmd, int (RespCommand.Set, (byte)SetOperation.SPOP) => SetPop(count, ptr, ref storageApi), (RespCommand.Set, (byte)SetOperation.SRANDMEMBER) => SetRandomMember(count, ptr, ref storageApi), (RespCommand.Set, (byte)SetOperation.SSCAN) => ObjectScan(count, ptr, GarnetObjectType.Set, ref storageApi), + (RespCommand.Set, (byte)SetOperation.SUNION) => SetUnion(count, ptr, ref storageApi), _ => ProcessOtherCommands(cmd, subcmd, count, ref storageApi), }; return success; diff --git a/libs/server/Storage/Session/ObjectStore/SetOps.cs b/libs/server/Storage/Session/ObjectStore/SetOps.cs index bcb6aba780..c88ce54611 100644 --- a/libs/server/Storage/Session/ObjectStore/SetOps.cs +++ b/libs/server/Storage/Session/ObjectStore/SetOps.cs @@ -2,6 +2,8 @@ // Licensed under the MIT license. using System; +using System.Collections.Generic; +using System.Diagnostics; using System.Text; using Garnet.common; using Tsavorite.core; @@ -361,6 +363,58 @@ public unsafe GarnetStatus SetScan(ArgSlice key, long cursor, st } + /// + /// Returns the members of the set resulting from the union of all the given sets. + /// Keys that do not exist are considered to be empty sets. + /// + /// + /// + /// + /// + /// + public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output, ref TObjectContext objectStoreContext) + where TObjectContext : ITsavoriteContext + { + output = new HashSet(new ByteArrayComparer()); + + if (keys.Length == 0) + return GarnetStatus.OK; + + var createTransaction = false; + + if (txnManager.state != TxnState.Running) + { + Debug.Assert(txnManager.state == TxnState.None); + createTransaction = true; + foreach (var item in keys) + txnManager.SaveKeyEntryToLock(item, true, LockType.Shared); + _ = txnManager.Run(true); + } + + // SetObject + var setObjectStoreLockableContext = txnManager.ObjectStoreLockableContext; + + try + { + foreach (var key in keys) + { + if (GET(key.ToArray(), out var currObject, ref setObjectStoreLockableContext) == GarnetStatus.OK) + { + var currSet = ((SetObject)currObject.garnetObject).Set; + output.UnionWith(currSet); + } + } + } + finally + { + if (createTransaction) + txnManager.Commit(true); + } + + return GarnetStatus.OK; + } + + /// /// Adds the specified members to the set at key. /// Specified members that are already a member of this set are ignored. diff --git a/libs/server/Transaction/TxnKeyManager.cs b/libs/server/Transaction/TxnKeyManager.cs index 1218e4c535..0034faa704 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -48,7 +48,7 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan RespCommand.SortedSet => SortedSetObjectKeys(subCommand, inputCount), RespCommand.List => ListObjectKeys(subCommand), RespCommand.Hash => HashObjectKeys(subCommand), - RespCommand.Set => SetObjectKeys(subCommand), + RespCommand.Set => SetObjectKeys(subCommand, inputCount), RespCommand.GET => SingleKey(1, false, LockType.Shared), RespCommand.SET => SingleKey(1, false, LockType.Exclusive), RespCommand.GETRANGE => SingleKey(1, false, LockType.Shared), @@ -172,7 +172,7 @@ private int HashObjectKeys(byte subCommand) }; } - private int SetObjectKeys(byte subCommand) + private int SetObjectKeys(byte subCommand, int inputCount) { return subCommand switch { @@ -183,6 +183,7 @@ private int SetObjectKeys(byte subCommand) (byte)SetOperation.SRANDMEMBER => SingleKey(1, true, LockType.Exclusive), (byte)SetOperation.SPOP => SingleKey(1, true, LockType.Exclusive), (byte)SetOperation.SISMEMBER => SingleKey(1, true, LockType.Shared), + (byte)SetOperation.SUNION => ListKeys(inputCount, true, LockType.Shared), _ => -1 }; } diff --git a/test/Garnet.test/RespSetTest.cs b/test/Garnet.test/RespSetTest.cs index 4371b72cd4..bda852bcaa 100644 --- a/test/Garnet.test/RespSetTest.cs +++ b/test/Garnet.test/RespSetTest.cs @@ -8,6 +8,7 @@ using Garnet.server; using NUnit.Framework; using StackExchange.Redis; +using SetOperation = StackExchange.Redis.SetOperation; namespace Garnet.test { @@ -273,6 +274,52 @@ public void CanDoSScanWithCursor() Assert.AreEqual(l, $"value:{setEntries.Length - 1}"); } + [Test] + public void CanDoSetUnion() + { + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(0); + var redisValues1 = new RedisValue[] { "item-a", "item-b", "item-c", "item-d" }; + var result = db.SetAdd(new RedisKey("key1"), redisValues1); + Assert.AreEqual(4, result); + + result = db.SetAdd(new RedisKey("key2"), new RedisValue[] { "item-c" }); + Assert.AreEqual(1, result); + + result = db.SetAdd(new RedisKey("key3"), new RedisValue[] { "item-a", "item-c", "item-e" }); + Assert.AreEqual(3, result); + + var members = db.SetCombine(SetOperation.Union, new RedisKey[] { "key1", "key2", "key3" }); + RedisValue[] entries = new RedisValue[] { "item-a", "item-b", "item-c", "item-d", "item-e" }; + Assert.AreEqual(5, members.Length); + // assert two arrays are equal ignoring order + Assert.IsTrue(members.OrderBy(x => x).SequenceEqual(entries.OrderBy(x => x))); + + members = db.SetCombine(SetOperation.Union, new RedisKey[] { "key1", "key2", "key3", "_not_exists" }); + Assert.AreEqual(5, members.Length); + Assert.IsTrue(members.OrderBy(x => x).SequenceEqual(entries.OrderBy(x => x))); + + members = db.SetCombine(SetOperation.Union, new RedisKey[] { "_not_exists_1", "_not_exists_2", "_not_exists_3" }); + Assert.IsEmpty(members); + + members = db.SetCombine(SetOperation.Union, new RedisKey[] { "_not_exists_1", "key1", "_not_exists_2", "_not_exists_3" }); + Assert.AreEqual(4, members.Length); + Assert.IsTrue(members.OrderBy(x => x).SequenceEqual(redisValues1.OrderBy(x => x))); + + members = db.SetCombine(SetOperation.Union, new RedisKey[] { "key1", "key2" }); + Assert.AreEqual(4, members.Length); + Assert.IsTrue(members.OrderBy(x => x).SequenceEqual(redisValues1.OrderBy(x => x))); + + try + { + db.SetCombine(SetOperation.Union, new RedisKey[] { }); + Assert.Fail(); + } + catch (RedisServerException e) + { + Assert.AreEqual(string.Format(CmdStrings.GenericErrWrongNumArgs, "SUNION"), e.Message); + } + } #endregion @@ -303,7 +350,6 @@ public void CanAddAndListMembersLC() expectedResponse = "*2\r\n$7\r\n\"Hello\"\r\n$7\r\n\"World\"\r\n"; strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); Assert.AreEqual(expectedResponse, strResponse); - } [Test] @@ -570,6 +616,47 @@ public void MultiWithNonExistingSet() Assert.AreEqual(res.AsSpan().Slice(0, expectedResponse.Length).ToArray(), expectedResponse); } + [Test] + public void CanDoSetUnionLC() + { + using var lightClientRequest = TestUtils.CreateRequest(); + var response = lightClientRequest.SendCommand("SADD myset ItemOne ItemTwo ItemThree ItemFour"); + var expectedResponse = ":4\r\n"; + var strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + response = lightClientRequest.SendCommand("SUNION myset another_set", 5); + expectedResponse = "*4\r\n"; + strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + lightClientRequest.SendCommand("SADD another_set ItemOne ItemFive ItemTwo ItemSix ItemSeven"); + response = lightClientRequest.SendCommand("SUNION myset another_set", 8); + expectedResponse = "*7\r\n"; + strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + response = lightClientRequest.SendCommand("SUNION myset no_exist_set", 5); + expectedResponse = "*4\r\n"; + strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + response = lightClientRequest.SendCommand("SUNION no_exist_set myset no_exist_set another_set", 8); + expectedResponse = "*7\r\n"; + strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + response = lightClientRequest.SendCommand("SUNION myset", 5); + expectedResponse = "*4\r\n"; + strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + response = lightClientRequest.SendCommand("SUNION"); + expectedResponse = $"-{string.Format(CmdStrings.GenericErrWrongNumArgs, "SUNION")}\r\n"; + strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + } + #endregion diff --git a/website/docs/commands/api-compatibility.md b/website/docs/commands/api-compatibility.md index a2a8183ae4..4cf7d35d84 100644 --- a/website/docs/commands/api-compatibility.md +++ b/website/docs/commands/api-compatibility.md @@ -239,7 +239,7 @@ Note that this list is subject to change as we continue to expand our API comman | | [SRANDMEMBER](data-structures.md#srandmember) | ➕ | | | | [SREM](data-structures.md#srem) | ➕ | | | | [SSCAN](data-structures.md#sscan) | ➕ | | -| | SUNION | ➖ | | +| | [SUNION](data-structures.md#sunion) | ➕ | | | | SUNIONSTORE | ➖ | | | **SORTED SET** | BZPOP | ➖ | | | | BZPOPMAX | ➖ | | diff --git a/website/docs/commands/data-structures.md b/website/docs/commands/data-structures.md index 7e15159f2f..edd30493df 100644 --- a/website/docs/commands/data-structures.md +++ b/website/docs/commands/data-structures.md @@ -509,6 +509,19 @@ The **match** parameter allows to apply a filter to elements after they have bee --- +### SUNION + +#### Syntax + +```bash + SUNION key [key ...] +``` + +Returns the members of the set resulting from the union of all the given sets. +Keys that do not exist are considered to be empty sets. + +--- + ## Sorted Set ### ZADD