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

Combining ZRange options in SortedSetObject #1021

Merged
merged 8 commits into from
Feb 18, 2025
Merged

Combining ZRange options in SortedSetObject #1021

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b0c07de
Combining ZRange options
Vijay-Nirmal Feb 14, 2025
1e9591f
Merge branch 'main' into chore/combining-zrange
Vijay-Nirmal Feb 14, 2025
be7c803
SortedSetRange ele fix
Vijay-Nirmal Feb 14, 2025
7b49dfe
Fixed test
Vijay-Nirmal Feb 14, 2025
ef0117d
Simplified rangeOpts creation and avoided some scratchBufferManager
Vijay-Nirmal Feb 15, 2025
e8942ea
Merge branch 'main' into chore/combining-zrange
Vijay-Nirmal Feb 15, 2025
331edcb
Merge branch 'main' into chore/combining-zrange
TalZaccai Feb 18, 2025
ee1c686
Merge branch 'main' into chore/combining-zrange
TalZaccai Feb 18, 2025
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
52 changes: 33 additions & 19 deletions libs/server/Objects/SortedSet/SortedSetObject.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,18 +27,12 @@ public enum SortedSetOperation : byte
ZINCRBY,
ZRANK,
ZRANGE,
ZRANGEBYLEX,
ZRANGEBYSCORE,
ZRANGESTORE,
GEOADD,
GEOHASH,
GEODIST,
GEOPOS,
GEOSEARCH,
GEOSEARCHSTORE,
ZREVRANGE,
ZREVRANGEBYLEX,
ZREVRANGEBYSCORE,
ZREVRANK,
ZREMRANGEBYLEX,
ZREMRANGEBYRANK,
Expand All @@ -51,6 +45,38 @@ public enum SortedSetOperation : byte
ZMSCORE
}

/// <summary>
/// Options for specifying the range in sorted set operations.
/// </summary>
[Flags]
public enum SortedSetRangeOpts : byte
{
/// <summary>
/// No options specified.
/// </summary>
None = 0,
/// <summary>
/// Range by score.
/// </summary>
ByScore = 1,
/// <summary>
/// Range by lexicographical order.
/// </summary>
ByLex = 1 << 1,
/// <summary>
/// Reverse the range order.
/// </summary>
Reverse = 1 << 2,
/// <summary>
/// Store the result.
/// </summary>
Store = 1 << 3,
/// <summary>
/// Include scores in the result.
/// </summary>
WithScores = 1 << 4
}

[Flags]
public enum SortedSetAddOption
{
Expand Down Expand Up @@ -259,14 +285,6 @@ public override unsafe bool Operate(ref ObjectInput input, ref GarnetObjectStore
case SortedSetOperation.ZRANK:
SortedSetRank(ref input, ref output.SpanByteAndMemory);
break;
case SortedSetOperation.ZRANGE:
case SortedSetOperation.ZRANGESTORE:
SortedSetRange(ref input, ref output.SpanByteAndMemory);
break;
case SortedSetOperation.ZRANGEBYLEX:
case SortedSetOperation.ZRANGEBYSCORE:
SortedSetRange(ref input, ref output.SpanByteAndMemory);
break;
case SortedSetOperation.GEOADD:
GeoAdd(ref input, ref output.SpanByteAndMemory);
break;
Expand All @@ -283,11 +301,7 @@ public override unsafe bool Operate(ref ObjectInput input, ref GarnetObjectStore
case SortedSetOperation.GEOSEARCHSTORE:
GeoSearch(ref input, ref output.SpanByteAndMemory);
break;
case SortedSetOperation.ZREVRANGE:
SortedSetRange(ref input, ref output.SpanByteAndMemory);
break;
case SortedSetOperation.ZREVRANGEBYLEX:
case SortedSetOperation.ZREVRANGEBYSCORE:
case SortedSetOperation.ZRANGE:
SortedSetRange(ref input, ref output.SpanByteAndMemory);
break;
case SortedSetOperation.ZREVRANK:
Expand Down
30 changes: 7 additions & 23 deletions libs/server/Objects/SortedSet/SortedSetObjectImpl.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -416,6 +416,7 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output)
//ZRANGE key min max [BYSCORE|BYLEX] [REV] [LIMIT offset count] [WITHSCORES]
//ZRANGEBYSCORE key min max [WITHSCORES] [LIMIT offset count]
//ZREVRANGEBYSCORE key max min [WITHSCORES] [LIMIT offset count]
var rangeOpts = (SortedSetRangeOpts)input.arg2;
var count = input.parseState.Count;
var respProtocolVersion = input.arg1;

Expand All @@ -436,30 +437,13 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output)
var maxSpan = input.parseState.GetArgSliceByRef(currIdx++).ReadOnlySpan;

// read the rest of the arguments
ZRangeOptions options = new();
switch (input.header.SortedSetOp)
ZRangeOptions options = new()
{
case SortedSetOperation.ZRANGEBYLEX:
options.ByLex = true;
break;
case SortedSetOperation.ZRANGESTORE:
options.WithScores = true;
break;
case SortedSetOperation.ZRANGEBYSCORE:
options.ByScore = true;
break;
case SortedSetOperation.ZREVRANGE:
options.Reverse = true;
break;
case SortedSetOperation.ZREVRANGEBYSCORE:
options.ByScore = true;
options.Reverse = true;
break;
case SortedSetOperation.ZREVRANGEBYLEX:
options.ByLex = true;
options.Reverse = true;
break;
}
ByScore = (rangeOpts & SortedSetRangeOpts.ByScore) != 0,
ByLex = (rangeOpts & SortedSetRangeOpts.ByLex) != 0,
Reverse = (rangeOpts & SortedSetRangeOpts.Reverse) != 0,
WithScores = (rangeOpts & SortedSetRangeOpts.WithScores) != 0 || (rangeOpts & SortedSetRangeOpts.Store) != 0
};

if (count > 2)
{
Expand Down
44 changes: 29 additions & 15 deletions libs/server/Resp/Objects/SortedSetCommands.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,20 +159,34 @@ private unsafe bool SortedSetRange<TGarnetApi>(RespCommand command, ref TGarnetA
var sbKey = parseState.GetArgSliceByRef(0).SpanByte;
var keyBytes = sbKey.ToByteArray();

var op =
command switch
{
RespCommand.ZRANGE => SortedSetOperation.ZRANGE,
RespCommand.ZREVRANGE => SortedSetOperation.ZREVRANGE,
RespCommand.ZRANGEBYLEX => SortedSetOperation.ZRANGEBYLEX,
RespCommand.ZRANGEBYSCORE => SortedSetOperation.ZRANGEBYSCORE,
RespCommand.ZREVRANGEBYLEX => SortedSetOperation.ZREVRANGEBYLEX,
RespCommand.ZREVRANGEBYSCORE => SortedSetOperation.ZREVRANGEBYSCORE,
_ => throw new Exception($"Unexpected {nameof(SortedSetOperation)}: {command}")
};
var rangeOpts = SortedSetRangeOpts.None;

var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = op };
var input = new ObjectInput(header, ref parseState, startIdx: 1, arg1: respProtocolVersion);
switch (command)
{
case RespCommand.ZRANGE:
break;
case RespCommand.ZREVRANGE:
rangeOpts = SortedSetRangeOpts.Reverse;
break;
case RespCommand.ZRANGEBYLEX:
rangeOpts = SortedSetRangeOpts.ByLex;
break;
case RespCommand.ZRANGEBYSCORE:
rangeOpts = SortedSetRangeOpts.ByScore;
break;
case RespCommand.ZREVRANGEBYLEX:
rangeOpts = SortedSetRangeOpts.ByLex | SortedSetRangeOpts.Reverse;
break;
case RespCommand.ZREVRANGEBYSCORE:
rangeOpts = SortedSetRangeOpts.ByScore | SortedSetRangeOpts.Reverse;
break;
case RespCommand.ZRANGESTORE:
rangeOpts = SortedSetRangeOpts.Store;
break;
}

var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE };
var input = new ObjectInput(header, ref parseState, startIdx: 1, arg1: respProtocolVersion, arg2: (int)rangeOpts);

var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(dcurr, (int)(dend - dcurr)) };

Expand Down Expand Up @@ -208,8 +222,8 @@ private unsafe bool SortedSetRangeStore<TGarnetApi>(ref TGarnetApi storageApi)
var dstKey = parseState.GetArgSliceByRef(0);
var srcKey = parseState.GetArgSliceByRef(1);

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

var status = storageApi.SortedSetRangeStore(dstKey, srcKey, ref input, out int result);

Expand Down
37 changes: 9 additions & 28 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 @@ -466,37 +466,18 @@ public unsafe GarnetStatus SortedSetRange<TObjectContext>(ArgSlice key, ArgSlice
return GarnetStatus.NOTFOUND;
}

ReadOnlySpan<byte> operation = default;
var sortedOperation = SortedSetOperation.ZRANGE;
switch (sortedSetOrderOperation)
var rangeOpts = sortedSetOrderOperation switch
{
case SortedSetOrderOperation.ByScore:
sortedOperation = SortedSetOperation.ZRANGEBYSCORE;
operation = "BYSCORE"u8;
break;
case SortedSetOrderOperation.ByLex:
sortedOperation = SortedSetOperation.ZRANGE;
operation = "BYLEX"u8;
break;
case SortedSetOrderOperation.ByRank:
if (reverse)
sortedOperation = SortedSetOperation.ZREVRANGE;
operation = default;
break;
}
SortedSetOrderOperation.ByScore => SortedSetRangeOpts.ByScore,
SortedSetOrderOperation.ByLex => SortedSetRangeOpts.ByLex,
_ => SortedSetRangeOpts.None
};

var arguments = new List<ArgSlice> { min, max };

// Operation order
if (!operation.IsEmpty)
{
arguments.Add(scratchBufferManager.CreateArgSlice(operation));
}

// Reverse
if (sortedOperation != SortedSetOperation.ZREVRANGE && reverse)
if (reverse)
{
arguments.Add(scratchBufferManager.CreateArgSlice("REV"u8));
rangeOpts |= SortedSetRangeOpts.Reverse;
}

// Limit parameter
Expand All @@ -517,9 +498,9 @@ public unsafe GarnetStatus SortedSetRange<TObjectContext>(ArgSlice key, ArgSlice
parseState.InitializeWithArguments([.. arguments]);

// Prepare the input
var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = sortedOperation };
var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE };
var inputArg = 2; // Default RESP server protocol version
var input = new ObjectInput(header, ref parseState, arg1: inputArg);
var input = new ObjectInput(header, ref parseState, arg1: inputArg, arg2: (int)rangeOpts);

var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(null) };
var status = ReadObjectStoreOperationWithOutput(key.ToArray(), ref input, ref objectContext, ref outputFooter);
Expand Down
14 changes: 5 additions & 9 deletions libs/server/Transaction/TxnKeyManager.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan<byte>
RespCommand.ZINCRBY => SortedSetObjectKeys(SortedSetOperation.ZINCRBY, inputCount),
RespCommand.ZRANK => SortedSetObjectKeys(SortedSetOperation.ZRANK, inputCount),
RespCommand.ZRANGE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
RespCommand.ZRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGEBYLEX, inputCount),
RespCommand.ZRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGEBYSCORE, inputCount),
RespCommand.ZRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
RespCommand.ZRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
RespCommand.ZREVRANK => SortedSetObjectKeys(SortedSetOperation.ZREVRANK, inputCount),
RespCommand.ZREMRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZREMRANGEBYLEX, inputCount),
RespCommand.ZREMRANGEBYRANK => SortedSetObjectKeys(SortedSetOperation.ZREMRANGEBYRANK, inputCount),
Expand All @@ -105,9 +105,9 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan<byte>
RespCommand.GEODIST => SortedSetObjectKeys(SortedSetOperation.GEODIST, inputCount),
RespCommand.GEOPOS => SortedSetObjectKeys(SortedSetOperation.GEOPOS, inputCount),
RespCommand.GEOSEARCH => SortedSetObjectKeys(SortedSetOperation.GEOSEARCH, inputCount),
RespCommand.ZREVRANGE => SortedSetObjectKeys(SortedSetOperation.ZREVRANGE, inputCount),
RespCommand.ZREVRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZREVRANGEBYLEX, inputCount),
RespCommand.ZREVRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZREVRANGEBYSCORE, inputCount),
RespCommand.ZREVRANGE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
RespCommand.ZREVRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
RespCommand.ZREVRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
RespCommand.LINDEX => ListObjectKeys((byte)ListOperation.LINDEX),
RespCommand.LINSERT => ListObjectKeys((byte)ListOperation.LINSERT),
RespCommand.LLEN => ListObjectKeys((byte)ListOperation.LLEN),
Expand Down Expand Up @@ -210,7 +210,6 @@ private int SortedSetObjectKeys(SortedSetOperation command, int inputCount)
SortedSetOperation.ZINCRBY => SingleKey(1, true, LockType.Exclusive),
SortedSetOperation.ZRANK => SingleKey(1, true, LockType.Exclusive),
SortedSetOperation.ZRANGE => SingleKey(1, true, LockType.Shared),
SortedSetOperation.ZRANGEBYSCORE => SingleKey(1, true, LockType.Shared),
SortedSetOperation.ZREVRANK => SingleKey(1, true, LockType.Exclusive),
SortedSetOperation.ZREMRANGEBYLEX => SingleKey(1, true, LockType.Exclusive),
SortedSetOperation.ZREMRANGEBYRANK => SingleKey(1, true, LockType.Exclusive),
Expand All @@ -224,9 +223,6 @@ private int SortedSetObjectKeys(SortedSetOperation command, int inputCount)
SortedSetOperation.GEODIST => SingleKey(1, true, LockType.Shared),
SortedSetOperation.GEOPOS => SingleKey(1, true, LockType.Shared),
SortedSetOperation.GEOSEARCH => SingleKey(1, true, LockType.Shared),
SortedSetOperation.ZREVRANGE => SingleKey(1, true, LockType.Shared),
SortedSetOperation.ZREVRANGEBYLEX => SingleKey(1, true, LockType.Shared),
SortedSetOperation.ZREVRANGEBYSCORE => SingleKey(1, true, LockType.Shared),
_ => -1
};
}
Expand Down