diff --git a/libs/server/API/GarnetApiObjectCommands.cs b/libs/server/API/GarnetApiObjectCommands.cs index b224866b31..b684dba249 100644 --- a/libs/server/API/GarnetApiObjectCommands.cs +++ b/libs/server/API/GarnetApiObjectCommands.cs @@ -309,7 +309,11 @@ public GarnetStatus SetMove(ArgSlice sourceKey, ArgSlice destinationKey, ArgSlic => storageSession.SetMove(sourceKey, destinationKey, member, out smoveResult); public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output) - => storageSession.SetUnion(keys, out output, ref objectContext); + => storageSession.SetUnion(keys, out output); + + /// + public GarnetStatus SetUnionStore(byte[] key, ArgSlice[] keys, out int count) + => storageSession.SetUnionStore(key, keys, out count); /// public GarnetStatus SetDiff(ArgSlice[] keys, out HashSet members) diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index 21bda389cf..0b98827f65 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -546,6 +546,16 @@ public interface IGarnetApi : IGarnetReadApi, IGarnetAdvancedApi /// GarnetStatus SetRandomMember(byte[] key, ArgSlice input, ref GarnetObjectStoreOutput outputFooter); + /// + /// This command is equal to SUNION, but instead of returning the resulting set, it is stored in destination. + /// If destination already exists, it is overwritten. + /// + /// + /// + /// + /// + GarnetStatus SetUnionStore(byte[] key, ArgSlice[] keys, out int count); + /// /// This command is equal to SDIFF, but instead of returning the resulting set, it is stored in destination. /// If destination already exists, it is overwritten. diff --git a/libs/server/Objects/Set/SetObject.cs b/libs/server/Objects/Set/SetObject.cs index a7338a824a..5a6c879283 100644 --- a/libs/server/Objects/Set/SetObject.cs +++ b/libs/server/Objects/Set/SetObject.cs @@ -27,6 +27,7 @@ public enum SetOperation : byte SRANDMEMBER, SISMEMBER, SUNION, + SUNIONSTORE, SDIFF, SDIFFSTORE, } diff --git a/libs/server/Resp/Objects/SetCommands.cs b/libs/server/Resp/Objects/SetCommands.cs index 086b3a5ace..e56cfe9976 100644 --- a/libs/server/Resp/Objects/SetCommands.cs +++ b/libs/server/Resp/Objects/SetCommands.cs @@ -147,6 +147,64 @@ private bool SetUnion(int count, byte* ptr, ref TGarnetApi storageAp return true; } + /// + /// This command is equal to SUNION, but instead of returning the resulting set, it is stored in destination. + /// If destination already exists, it is overwritten. + /// + /// + /// + /// + /// + /// + private bool SetUnionStore(int count, byte* ptr, ref TGarnetApi storageApi) + where TGarnetApi : IGarnetApi + { + if (count < 2) + { + return AbortWithWrongNumberOfArguments("SUNIONSTORE", count); + } + + // Get the key + if (!RespReadUtils.ReadByteArrayWithLengthHeader(out var key, ref ptr, recvBufferPtr + bytesRead)) + return false; + + if (NetworkSingleKeySlotVerify(key, false)) + { + var bufSpan = new ReadOnlySpan(recvBufferPtr, bytesRead); + if (!DrainCommands(bufSpan, count)) + return false; + return true; + } + + var keys = new ArgSlice[count - 1]; + for (var i = 0; i < count - 1; 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; + } + + var status = storageApi.SetUnionStore(key, keys, out var output); + + if (status == GarnetStatus.OK) + { + while (!RespWriteUtils.WriteInteger(output, ref dcurr, dend)) + SendAndReset(); + } + + // Move input head + 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 3f24c000d7..87e6df1619 100644 --- a/libs/server/Resp/RespCommand.cs +++ b/libs/server/Resp/RespCommand.cs @@ -908,6 +908,10 @@ static RespCommand MatchedNone(RespServerSession session, int oldReadHead) { return (RespCommand.Set, (byte)SetOperation.SRANDMEMBER); } + else if (*(ulong*)(ptr + 2) == MemoryMarshal.Read("1\r\nSUNIO"u8) && *(ulong*)(ptr + 10) == MemoryMarshal.Read("NSTORE\r\n"u8)) + { + return (RespCommand.Set, (byte)SetOperation.SUNIONSTORE); + } break; case 12: diff --git a/libs/server/Resp/RespCommandsInfo.cs b/libs/server/Resp/RespCommandsInfo.cs index 62115b8a3e..a328328f58 100644 --- a/libs/server/Resp/RespCommandsInfo.cs +++ b/libs/server/Resp/RespCommandsInfo.cs @@ -212,18 +212,19 @@ public static RespCommandsInfo findCommand(RespCommand cmd, byte subCmd = 0) private static readonly Dictionary setCommandsInfoMap = new Dictionary { - {(byte)SetOperation.SADD, new RespCommandsInfo("SADD", RespCommand.Set, -2, null, (byte)SetOperation.SADD)}, - {(byte)SetOperation.SMEMBERS, new RespCommandsInfo("SMEMBERS", RespCommand.Set, 1, null, (byte)SetOperation.SMEMBERS)}, - {(byte)SetOperation.SREM, new RespCommandsInfo("SREM", RespCommand.Set, -2, null, (byte)SetOperation.SREM)}, - {(byte)SetOperation.SCARD, new RespCommandsInfo("SCARD", RespCommand.Set, 1, null, (byte)SetOperation.SCARD)}, - {(byte)SetOperation.SRANDMEMBER,new RespCommandsInfo("SRANDMEMBER", RespCommand.Set, -2, null, (byte)SetOperation.SRANDMEMBER)}, - {(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.SMOVE, new RespCommandsInfo("SMOVE", RespCommand.Set, 3, null, (byte)SetOperation.SMOVE) }, - {(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) }, - {(byte)SetOperation.SDIFF, new RespCommandsInfo("SDIFF", RespCommand.Set, -1, null, (byte)SetOperation.SDIFF) }, - {(byte)SetOperation.SDIFFSTORE, new RespCommandsInfo("SDIFFSTORE", RespCommand.Set, -2, null, (byte)SetOperation.SDIFFSTORE) } + {(byte)SetOperation.SADD, new RespCommandsInfo("SADD", RespCommand.Set, -2, null, (byte)SetOperation.SADD)}, + {(byte)SetOperation.SMEMBERS, new RespCommandsInfo("SMEMBERS", RespCommand.Set, 1, null, (byte)SetOperation.SMEMBERS)}, + {(byte)SetOperation.SREM, new RespCommandsInfo("SREM", RespCommand.Set, -2, null, (byte)SetOperation.SREM)}, + {(byte)SetOperation.SCARD, new RespCommandsInfo("SCARD", RespCommand.Set, 1, null, (byte)SetOperation.SCARD)}, + {(byte)SetOperation.SRANDMEMBER, new RespCommandsInfo("SRANDMEMBER", RespCommand.Set, -2, null, (byte)SetOperation.SRANDMEMBER)}, + {(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.SMOVE, new RespCommandsInfo("SMOVE", RespCommand.Set, 3, null, (byte)SetOperation.SMOVE) }, + {(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) }, + {(byte)SetOperation.SUNIONSTORE, new RespCommandsInfo("SUNIONSTORE", RespCommand.Set, -2, null, (byte)SetOperation.SUNIONSTORE) }, + {(byte)SetOperation.SDIFF, new RespCommandsInfo("SDIFF", RespCommand.Set, -1, null, (byte)SetOperation.SDIFF) }, + {(byte)SetOperation.SDIFFSTORE, new RespCommandsInfo("SDIFFSTORE", RespCommand.Set, -2, null, (byte)SetOperation.SDIFFSTORE) } }; private static readonly Dictionary customCommandsInfoMap = new Dictionary diff --git a/libs/server/Resp/RespInfo.cs b/libs/server/Resp/RespInfo.cs index 3c2f185b05..98f0d22bc0 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", "SUNION", "SDIFF", "SDIFFSTORE", "SMOVE", + "SADD", "SREM", "SPOP", "SMEMBERS", "SCARD", "SSCAN", "SRANDMEMBER", "SISMEMBER", "SUNION", "SUNIONSTORE", "SDIFF", "SDIFFSTORE", "SMOVE", //Scan ops "DBSIZE", "KEYS","SCAN", // Geospatial commands diff --git a/libs/server/Resp/RespServerSession.cs b/libs/server/Resp/RespServerSession.cs index f0f53fca10..9d6cf31831 100644 --- a/libs/server/Resp/RespServerSession.cs +++ b/libs/server/Resp/RespServerSession.cs @@ -510,6 +510,7 @@ private bool ProcessArrayCommands(RespCommand cmd, byte subcmd, int (RespCommand.Set, (byte)SetOperation.SSCAN) => ObjectScan(count, ptr, GarnetObjectType.Set, ref storageApi), (RespCommand.Set, (byte)SetOperation.SMOVE) => SetMove(count, ptr, ref storageApi), (RespCommand.Set, (byte)SetOperation.SUNION) => SetUnion(count, ptr, ref storageApi), + (RespCommand.Set, (byte)SetOperation.SUNIONSTORE) => SetUnionStore(count, ptr, ref storageApi), (RespCommand.Set, (byte)SetOperation.SDIFF) => SetDiff(count, ptr, ref storageApi), (RespCommand.Set, (byte)SetOperation.SDIFFSTORE) => SetDiffStore(count, ptr, ref storageApi), _ => ProcessOtherCommands(cmd, subcmd, count, ref storageApi), diff --git a/libs/server/Storage/Session/ObjectStore/SetOps.cs b/libs/server/Storage/Session/ObjectStore/SetOps.cs index a9736e7de0..2bd7c11061 100644 --- a/libs/server/Storage/Session/ObjectStore/SetOps.cs +++ b/libs/server/Storage/Session/ObjectStore/SetOps.cs @@ -428,11 +428,8 @@ internal unsafe GarnetStatus SetMove(ArgSlice sourceKey, ArgSlice destinationKey /// /// /// - /// - /// /// - public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output, ref TObjectContext objectStoreContext) - where TObjectContext : ITsavoriteContext + public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet output) { output = new HashSet(new ByteArrayComparer()); @@ -455,14 +452,58 @@ public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet + /// This command is equal to SUNION, but instead of returning the resulting set, it is stored in destination. + /// If destination already exists, it is overwritten. + /// + /// + /// + /// + /// + public GarnetStatus SetUnionStore(byte[] key, ArgSlice[] keys, out int count) + { + count = default; + + var destination = scratchBufferManager.CreateArgSlice(key); + + var createTransaction = false; + + if (txnManager.state != TxnState.Running) + { + Debug.Assert(txnManager.state == TxnState.None); + createTransaction = true; + txnManager.SaveKeyEntryToLock(destination, true, LockType.Exclusive); + foreach (var item in keys) + txnManager.SaveKeyEntryToLock(item, true, LockType.Shared); + _ = txnManager.Run(true); + } + + // SetObject + var setObjectStoreLockableContext = txnManager.ObjectStoreLockableContext; + + try + { + var members = SetUnion(keys, ref setObjectStoreLockableContext); + + var newSetObject = new SetObject(); + foreach (var item in members) { - if (GET(key.ToArray(), out var currObject, ref setObjectStoreLockableContext) == GarnetStatus.OK) - { - var currSet = ((SetObject)currObject.garnetObject).Set; - output.UnionWith(currSet); - } + _ = newSetObject.Set.Add(item); + newSetObject.UpdateSize(item); } + _ = SET(key, newSetObject, ref setObjectStoreLockableContext); + count = members.Count; } finally { @@ -473,6 +514,27 @@ public GarnetStatus SetUnion(ArgSlice[] keys, out HashSet SetUnion(ArgSlice[] keys, ref TObjectContext objectContext) + where TObjectContext : ITsavoriteContext + { + var result = new HashSet(new ByteArrayComparer()); + if (keys.Length == 0) + { + return result; + } + + foreach (var item in keys) + { + if (GET(item.ToArray(), out var currObject, ref objectContext) == GarnetStatus.OK) + { + var currSet = ((SetObject)currObject.garnetObject).Set; + result.UnionWith(currSet); + } + } + + return result; + } + /// /// 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 976234791b..c1c19116c0 100644 --- a/libs/server/Transaction/TxnKeyManager.cs +++ b/libs/server/Transaction/TxnKeyManager.cs @@ -184,6 +184,7 @@ private int SetObjectKeys(byte subCommand, int inputCount) (byte)SetOperation.SPOP => SingleKey(1, true, LockType.Exclusive), (byte)SetOperation.SISMEMBER => SingleKey(1, true, LockType.Shared), (byte)SetOperation.SUNION => ListKeys(inputCount, true, LockType.Shared), + (byte)SetOperation.SUNIONSTORE => XSTOREKeys(inputCount, true), (byte)SetOperation.SDIFF => ListKeys(inputCount, true, LockType.Shared), (byte)SetOperation.SDIFFSTORE => XSTOREKeys(inputCount, true), (byte)SetOperation.SMOVE => ListKeys(inputCount, true, LockType.Exclusive), diff --git a/test/Garnet.test/RespSetTest.cs b/test/Garnet.test/RespSetTest.cs index 3c9bb65067..4412ac1547 100644 --- a/test/Garnet.test/RespSetTest.cs +++ b/test/Garnet.test/RespSetTest.cs @@ -322,6 +322,35 @@ public void CanDoSetUnion() } } + [Test] + [TestCase("key")] + [TestCase("")] + public void CanDoSetUnionStore(string key) + { + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(0); + + var key1 = "key1"; + var key1Value = new RedisValue[] { "a", "b", "c" }; + + var key2 = "key2"; + var key2Value = new RedisValue[] { "c", "d", "e" }; + + var addResult = db.SetAdd(key1, key1Value); + Assert.AreEqual(3, addResult); + addResult = db.SetAdd(key2, key2Value); + Assert.AreEqual(3, addResult); + + var result = (int)db.Execute("SUNIONSTORE", key, key1, key2); + Assert.AreEqual(5, result); + + var membersResult = db.SetMembers(key); + Assert.AreEqual(5, membersResult.Length); + var strResult = membersResult.Select(m => m.ToString()).ToArray(); + var expectedResult = new[] { "a", "b", "c", "d", "e" }; + Assert.IsTrue(expectedResult.OrderBy(t => t).SequenceEqual(strResult.OrderBy(t => t))); + } + [Test] public void CanDoSdiff() { @@ -854,6 +883,21 @@ public void CanDoSetUnionLC() Assert.AreEqual(expectedResponse, strResponse); } + [Test] + public void CanDoSunionStoreLC() + { + var lightClientRequest = TestUtils.CreateRequest(); + _ = lightClientRequest.SendCommand("SADD key1 a b c"); + _ = lightClientRequest.SendCommand("SADD key2 c d e"); + var response = lightClientRequest.SendCommand("SUNIONSTORE key key1 key2"); + var expectedResponse = ":5\r\n"; + Assert.AreEqual(expectedResponse, response.AsSpan().Slice(0, expectedResponse.Length).ToArray()); + + var membersResponse = lightClientRequest.SendCommand("SMEMBERS key"); + expectedResponse = "*5\r\n$1\r\na\r\n$1\r\nb\r\n$1\r\nc\r\n$1\r\nd\r\n$1\r\ne\r\n"; + Assert.AreEqual(expectedResponse, membersResponse.AsSpan().Slice(0, expectedResponse.Length).ToArray()); + } + [Test] public void CanDoSdiffLC() { @@ -941,6 +985,21 @@ public void CanDoSdiffStoreWhenMemberKeysNotExisting() strResponse = Encoding.ASCII.GetString(membersResponse).Substring(0, expectedResponse.Length); Assert.AreEqual(expectedResponse, strResponse); } + + [Test] + public void CanDoSunionStoreWhenMemberKeysNotExisting() + { + using var lightClientRequest = TestUtils.CreateRequest(); + var response = lightClientRequest.SendCommand("SUNIONSTORE key key1 key2 key3"); + var expectedResponse = ":0\r\n"; + var strResponse = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length); + Assert.AreEqual(expectedResponse, strResponse); + + var membersResponse = lightClientRequest.SendCommand("SMEMBERS key"); + expectedResponse = "*0\r\n"; + strResponse = Encoding.ASCII.GetString(membersResponse).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 e76531e489..cc9c7f0ae1 100644 --- a/website/docs/commands/api-compatibility.md +++ b/website/docs/commands/api-compatibility.md @@ -240,7 +240,7 @@ Note that this list is subject to change as we continue to expand our API comman | | [SREM](data-structures.md#srem) | ➕ | | | | [SSCAN](data-structures.md#sscan) | ➕ | | | | [SUNION](data-structures.md#sunion) | ➕ | | -| | SUNIONSTORE | ➖ | | +| | [SUNIONSTORE](data-structures.md#sunionstore) | ➕ | | | **SORTED SET** | BZPOP | ➖ | | | | BZPOPMAX | ➖ | | | | BZPOPMIN | ➖ | | diff --git a/website/docs/commands/data-structures.md b/website/docs/commands/data-structures.md index 5749513312..7e0486818d 100644 --- a/website/docs/commands/data-structures.md +++ b/website/docs/commands/data-structures.md @@ -522,6 +522,20 @@ Keys that do not exist are considered to be empty sets. --- +### SUNIONSTORE + +#### Syntax + +```bash + SUNIONSTORE destination key [key ...] +``` + +This command is equal to [SUNION](#SUNION), but instead of returning the resulting set, it is stored in **destination**. + +If **destination** already exists, it is overwritten. + +--- + ### SDIFF #### Syntax