Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Compatibility] Added ZRANGESTORE command #826

Merged
merged 15 commits into from
Dec 8, 2024
Merged

[Compatibility] Added ZRANGESTORE command #826

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
e055f42
WIP of ZRANGESTORE
Vijay-Nirmal Nov 23, 2024
d26ca19
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Nov 24, 2024
926f278
Final changes
Vijay-Nirmal Nov 24, 2024
db7db2e
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Nov 24, 2024
aea3891
Added ACL test
Vijay-Nirmal Nov 24, 2024
2c65e26
Removed temp code
Vijay-Nirmal Nov 24, 2024
376b256
Removed TODO command
Vijay-Nirmal Nov 24, 2024
1d8216f
Slot Test
Vijay-Nirmal Nov 25, 2024
dd426e4
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Nov 26, 2024
a027004
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Nov 27, 2024
91915e0
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Nov 28, 2024
c86d1f8
Merge latest
Vijay-Nirmal Dec 5, 2024
20ac282
Review command fix
Vijay-Nirmal Dec 5, 2024
98e1b86
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Dec 7, 2024
4295c1e
Merge branch 'main' into new/ZRANGESTORE-command
Vijay-Nirmal Dec 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Final changes
  • Loading branch information
@Vijay-Nirmal
Vijay-Nirmal committed Nov 24, 2024
commit 926f278c1d6cf8911feb29b855bbb1fa70bacaae
2 changes: 2 additions & 0 deletions libs/server/Objects/SortedSet/SortedSetObject.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ public enum SortedSetOperation : byte
ZRANK,
ZRANGE,
ZRANGEBYSCORE,
ZRANGESTORE,
GEOADD,
GEOHASH,
GEODIST,
Expand Down Expand Up @@ -256,6 +257,7 @@ public override unsafe bool Operate(ref ObjectInput input, ref SpanByteAndMemory
SortedSetRank(ref input, ref output);
break;
case SortedSetOperation.ZRANGE:
case SortedSetOperation.ZRANGESTORE:
SortedSetRange(ref input, ref output);
break;
case SortedSetOperation.ZRANGEBYSCORE:
Expand Down
3 changes: 3 additions & 0 deletions libs/server/Objects/SortedSet/SortedSetObjectImpl.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -435,6 +435,9 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output)
ZRangeOptions options = new();
switch (input.header.SortedSetOp)
{
case SortedSetOperation.ZRANGESTORE:
options.WithScores = true;
break;
case SortedSetOperation.ZRANGEBYSCORE:
options.ByScore = true;
break;
Expand Down
4 changes: 2 additions & 2 deletions libs/server/Resp/Objects/SortedSetCommands.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,7 @@ private unsafe bool SortedSetRange<TGarnetApi>(RespCommand command, ref TGarnetA
return true;
}

private unsafe bool SortedSetRangeStore<TGarnetApi>(RespCommand command, ref TGarnetApi storageApi)
private unsafe bool SortedSetRangeStore<TGarnetApi>(ref TGarnetApi storageApi)
where TGarnetApi : IGarnetApi
{
// ZRANGESTORE dst src min max [BYSCORE | BYLEX] [REV] [LIMIT offset count]
Expand All @@ -205,7 +205,7 @@ private unsafe bool SortedSetRangeStore<TGarnetApi>(RespCommand command, ref TGa
var destKey = parseState.GetArgSliceByRef(0);
var sbKey = parseState.GetArgSliceByRef(1);

var header = new RespInputHeader(GarnetObjectType.SortedSet);
var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGESTORE };
var input = new ObjectInput(header, ref parseState, startIdx: 2, arg1: respProtocolVersion);

var status = storageApi.SortedSetRangeStore(destKey, sbKey, ref input, out int result);
Expand Down
2 changes: 1 addition & 1 deletion libs/server/Resp/RespServerSession.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -598,7 +598,7 @@ private bool ProcessArrayCommands<TGarnetApi>(RespCommand cmd, ref TGarnetApi st
RespCommand.ZINCRBY => SortedSetIncrement(ref storageApi),
RespCommand.ZRANK => SortedSetRank(cmd, ref storageApi),
RespCommand.ZRANGE => SortedSetRange(cmd, ref storageApi),
RespCommand.ZRANGE => SortedSetRangeStore(ref storageApi),
RespCommand.ZRANGESTORE => SortedSetRangeStore(ref storageApi),
RespCommand.ZRANGEBYSCORE => SortedSetRange(cmd, ref storageApi),
RespCommand.ZREVRANK => SortedSetRank(cmd, ref storageApi),
RespCommand.ZREMRANGEBYLEX => SortedSetLengthByValue(cmd, ref storageApi),
Expand Down
67 changes: 31 additions & 36 deletions libs/server/Storage/Session/ObjectStore/SortedSetOps.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -702,7 +702,7 @@ public GarnetStatus SortedSetAdd<TObjectContext>(byte[] key, ref ObjectInput inp
/// <param name="result">The result of the operation, indicating the number of elements stored.</param>
/// <param name="objectStoreContext">The context of the object store.</param>
/// <returns>Returns a GarnetStatus indicating the success or failure of the operation.</returns>
public GarnetStatus SortedSetRangeStore<TObjectContext>(ArgSlice distKey, ArgSlice sbKey, ref ObjectInput input, out int result, ref TObjectContext objectStoreContext)
public unsafe GarnetStatus SortedSetRangeStore<TObjectContext>(ArgSlice distKey, ArgSlice sbKey, ref ObjectInput input, out int result, ref TObjectContext objectStoreContext)
where TObjectContext : ITsavoriteContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions, ObjectStoreFunctions, ObjectStoreAllocator>
{
result = 0;
Expand All @@ -722,19 +722,14 @@ public GarnetStatus SortedSetRangeStore<TObjectContext>(ArgSlice distKey, ArgSli
}

// SetObject
var setObjectStoreLockableContext = txnManager.ObjectStoreLockableContext;
var objectStoreLockableContext = txnManager.ObjectStoreLockableContext;

try
{
var rangeParseState = new SessionParseState();
rangeParseState.Initialize(input.parseState.Count - 1);
rangeParseState.SetArguments(0, input.parseState.Parameters.Slice(1));
rangeParseState.SetArguments(input.parseState.Count - 2, ArgSlice.FromPinnedSpan(CmdStrings.WITHSCORES));

var rangeInput = new ObjectInput(input.header, ref rangeParseState);
SpanByteAndMemory rangeOutputMem = default;
var rangeOutput = new GarnetObjectStoreOutput() { spanByteAndMemory = rangeOutputMem };
var status = SortedSetRange(sbKey.ToArray(), ref rangeInput, ref rangeOutput, ref objectStoreContext);
var status = SortedSetRange(sbKey.ToArray(), ref input, ref rangeOutput, ref objectStoreLockableContext);
TalZaccai marked this conversation as resolved.
Show resolved Hide resolved
rangeOutputMem = rangeOutput.spanByteAndMemory;

if (status == GarnetStatus.WRONGTYPE)
{
Expand All @@ -751,45 +746,45 @@ public GarnetStatus SortedSetRangeStore<TObjectContext>(ArgSlice distKey, ArgSli
Debug.Assert(!rangeOutputMem.IsSpanByte, "Output should not be in SpanByte format when the status is OK");

var rangeOutputHandler = rangeOutputMem.Memory.Memory.Pin();

if (status == GarnetStatus.OK)
try
{
var rangeOutPtr = (byte*)rangeOutputHandler.Pointer;
ref var currOutPtr = ref rangeOutPtr;
var endOutPtr = rangeOutPtr + rangeOutputMem.Length;

var destinationKey = distKey.ToArray();
objectStoreLockableContext.Delete(ref destinationKey);

var zParseState = new SessionParseState();
zParseState.Initialize(foundItems * 2);
RespReadUtils.ReadUnsignedArrayLength(out var arrayLen, ref currOutPtr, endOutPtr);
Debug.Assert(arrayLen % 2 == 0, "Should always contain element and its score");
result = arrayLen / 2;

for (int j = 0; j < foundItems; j++)
if (result > 0)
{
RespReadUtils.ReadUnsignedArrayLength(out var innerLength, ref currOutPtr, endOutPtr);
Debug.Assert(innerLength == 2, "Should always has location and hash or distance");
parseState.Initialize(arrayLen); // 2 elements per pair (result * 2)

RespReadUtils.TrySliceWithLengthHeader(out var location, ref currOutPtr, endOutPtr);
if (storeDistIdx != -1)
for (int j = 0; j < result; j++)
{
RespReadUtils.ReadSpanWithLengthHeader(out var score, ref currOutPtr, endOutPtr);
zParseState.SetArgument(2 * j, ArgSlice.FromPinnedSpan(score));
zParseState.SetArgument((2 * j) + 1, ArgSlice.FromPinnedSpan(location));
// Read member/element into parse state
parseState.Read((2 * j) + 1, ref currOutPtr, endOutPtr);
// Read score into parse state
parseState.Read(2 * j, ref currOutPtr, endOutPtr);
}
else
{
RespReadUtils.ReadIntegerAsSpan(out var score, ref currOutPtr, endOutPtr);
zParseState.SetArgument(2 * j, ArgSlice.FromPinnedSpan(score));
zParseState.SetArgument((2 * j) + 1, ArgSlice.FromPinnedSpan(location));
}
}

var zAddInput = new ObjectInput(new RespInputHeader
{
type = GarnetObjectType.SortedSet,
SortedSetOp = SortedSetOperation.ZADD,
}, ref zParseState);
var zAddInput = new ObjectInput(new RespInputHeader
{
type = GarnetObjectType.SortedSet,
SortedSetOp = SortedSetOperation.ZADD,
}, ref parseState);

var zAddOutput = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) };
RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
var zAddOutput = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) };
RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
}
}
finally
{
rangeOutputHandler.Dispose();
}

return status;
}
finally
Expand Down
90 changes: 90 additions & 0 deletions test/Garnet.test/RespSortedSetTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1086,6 +1086,41 @@ public void CheckSortedSetDifferenceStoreWithNoMatchSE()
ClassicAssert.AreEqual(0, actualMembers.Length);
}

[Test]
[TestCase("user1:obj1", "user1:objA", new[] { "Hello", "World" }, new[] { 1.0, 2.0 }, new[] { "Hello", "World" }, new[] { 1.0, 2.0 })] // Normal case
[TestCase("user1:emptySet", "user1:objB", new string[] { }, new double[] { }, new string[] { }, new double[] { })] // Empty set
[TestCase("user1:nonExistingKey", "user1:objC", new string[] { }, new double[] { }, new string[] { }, new double[] { })] // Non-existing key
[TestCase("user1:obj2", "user1:objD", new[] { "Alpha", "Beta", "Gamma" }, new[] { 1.0, 2.0, 3.0 }, new[] { "Beta", "Gamma" }, new[] { 2.0, 3.0 }, -2, -1)] // Negative range
public void CheckSortedSetRangeStoreSE(string key, string destinationKey, string[] elements, double[] scores, string[] expectedElements, double[] expectedScores, int start = 0, int stop = -1)
TalZaccai marked this conversation as resolved.
Show resolved Hide resolved
TalZaccai marked this conversation as resolved.
Show resolved Hide resolved
{
int expectedCount = expectedElements.Length;

using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig());
var db = redis.GetDatabase(0);

var redisKey = new RedisKey(key);
TalZaccai marked this conversation as resolved.
Show resolved Hide resolved
var redisDestinationKey = new RedisKey(destinationKey);
var keyValues = elements.Zip(scores, (e, s) => new SortedSetEntry(e, s)).ToArray();

// Set up sorted set if elements are provided
if (keyValues.Length > 0)
{
db.SortedSetAdd(redisKey, keyValues);
}

var actualCount = db.SortedSetRangeAndStore(redisKey, redisDestinationKey, start, stop);
ClassicAssert.AreEqual(expectedCount, actualCount);

var actualMembers = db.SortedSetRangeByScoreWithScores(redisDestinationKey);
ClassicAssert.AreEqual(expectedCount, actualMembers.Length);

for (int i = 0; i < expectedCount; i++)
{
ClassicAssert.AreEqual(expectedElements[i], actualMembers[i].Element.ToString());
ClassicAssert.AreEqual(expectedScores[i], actualMembers[i].Score);
}
}

#endregion

#region LightClientTests
Expand Down Expand Up @@ -1210,6 +1245,61 @@ public void CanDoZRangeByIndexLC(int bytesSent)
ClassicAssert.AreEqual(expectedResponse, actualValue);
}

// ZRANGEBSTORE
[Test]
[TestCase("user1:obj1", "user1:objA", new[] { "Hello", "World" }, new[] { 1.0, 2.0 }, new[] { "Hello", "World" }, new[] { 1.0, 2.0 })] // Normal case
[TestCase("user1:emptySet", "user1:objB", new string[] { }, new double[] { }, new string[] { }, new double[] { })] // Empty set
[TestCase("user1:nonExistingKey", "user1:objC", new string[] { }, new double[] { }, new string[] { }, new double[] { })] // Non-existing key
[TestCase("user1:obj2", "user1:objD", new[] { "Alpha", "Beta", "Gamma" }, new[] { 1.0, 2.0, 3.0 }, new[] { "Beta", "Gamma" }, new[] { 2.0, 3.0 }, -2, -1)] // Negative range
public void CheckSortedSetRangeStoreLC(string key, string destinationKey, string[] elements, double[] scores, string[] expectedElements, double[] expectedScores, int start = 0, int stop = -1)
{
using var lightClientRequest = TestUtils.CreateRequest();

// Setup initial sorted set if elements exist
if (elements.Length > 0)
{
var addCommand = $"ZADD {key} " + string.Join(" ", elements.Zip(scores, (e, s) => $"{s} {e}"));
var response = lightClientRequest.SendCommand(addCommand);
var expectedResponse = $":{elements.Length}\r\n";
var actualValue = Encoding.ASCII.GetString(response).Substring(0, expectedResponse.Length);
ClassicAssert.AreEqual(expectedResponse, actualValue);
}

// Execute ZRANGESTORE
var rangeStoreCommand = $"ZRANGESTORE {destinationKey} {key} {start} {stop}";
var response2 = lightClientRequest.SendCommand(rangeStoreCommand);
var expectedResponse2 = $":{expectedElements.Length}\r\n";
var actualValue2 = Encoding.ASCII.GetString(response2).Substring(0, expectedResponse2.Length);
ClassicAssert.AreEqual(expectedResponse2, actualValue2);

// Verify stored result using ZRANGE
if (expectedElements.Length > 0)
{
var verifyCommand = $"ZRANGE {destinationKey} 0 -1 WITHSCORES";
var response3 = lightClientRequest.SendCommand(verifyCommand, expectedElements.Length * 2 + 1);
var expectedItems = new List<string>();
expectedItems.Add($"*{expectedElements.Length * 2}");
for (int i = 0; i < expectedElements.Length; i++)
{
expectedItems.Add($"${expectedElements[i].Length}");
expectedItems.Add(expectedElements[i]);
expectedItems.Add($"${expectedScores[i].ToString().Length}");
expectedItems.Add(expectedScores[i].ToString());
}
var expectedResponse3 = string.Join("\r\n", expectedItems) + "\r\n";
var actualValue3 = Encoding.ASCII.GetString(response3).Substring(0, expectedResponse3.Length);
ClassicAssert.AreEqual(expectedResponse3, actualValue3);
}
else
{
var verifyCommand = $"ZRANGE {destinationKey} 0 -1";
var response3 = lightClientRequest.SendCommand(verifyCommand);
var expectedResponse3 = "*0\r\n";
var actualValue3 = Encoding.ASCII.GetString(response3).Substring(0, expectedResponse3.Length);
ClassicAssert.AreEqual(expectedResponse3, actualValue3);
}
}

[Test]
[TestCase(10)]
[TestCase(50)]
Expand Down
2 changes: 1 addition & 1 deletion website/docs/commands/api-compatibility.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -339,7 +339,7 @@ Note that this list is subject to change as we continue to expand our API comman
| | [ZRANGE](data-structures.md#zrange) | ➕ | |
| | [ZRANGEBYLEX](data-structures.md#zrangebylex) | ➕ | (Deprecated) |
| | [ZRANGEBYSCORE](data-structures.md#zrangebyscore) | ➕ | (Deprecated) |
| | ZRANGESTORE | ➖ | |
| | [ZRANGESTORE](data-structures.md#zrangestore) | ➕ | |
| | [ZRANK](data-structures.md#zrank) | ➕ | |
| | [ZREM](data-structures.md#zrem) | ➕ | |
| | [ZREMRANGEBYLEX](data-structures.md#zremrangebylex) | ➕ | |
Expand Down
12 changes: 12 additions & 0 deletions website/docs/commands/data-structures.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1068,6 +1068,18 @@ If member does not exist in the sorted set, or **key** does not exist, nil is re

---

### ZRANGESTORE

#### Syntax

```bash
ZRANGESTORE dst src min max [BYSCORE|BYLEX] [REV] [LIMIT offset count]
```

Stores the specified range of elements in the sorted set stored at **src** into the sorted set stored at **dst**.

---

## Geospatial indices

### GEOADD
Expand Down
Loading