Skip to content

Commit

Permalink
Added some improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
@TalZaccai
TalZaccai committed Nov 8, 2024
1 parent 970d421 commit fa58e13
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 65 deletions.
2 changes: 2 additions & 0 deletions libs/server/Objects/SortedSet/SortedSetObject.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ public enum SortedSetOperation : byte
GEODIST,
GEOPOS,
GEOSEARCH,
GEOSEARCHSTORE,
ZREVRANGE,
ZREVRANGEBYSCORE,
ZREVRANK,
Expand Down Expand Up @@ -273,6 +274,7 @@ public override unsafe bool Operate(ref ObjectInput input, ref SpanByteAndMemory
GeoPosition(ref input, ref output);
break;
case SortedSetOperation.GEOSEARCH:
case SortedSetOperation.GEOSEARCHSTORE:
GeoSearch(ref input, ref output);
break;
case SortedSetOperation.ZREVRANGE:
Expand Down
7 changes: 4 additions & 3 deletions libs/server/Objects/SortedSetGeo/SortedSetGeoObjectImpl.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -396,9 +396,10 @@ private void GeoSearch(ref ObjectInput input, ref SpanByteAndMemory output)

opts.WithCount = true;
}
else if (tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHCOORD"u8)) opts.WithCoord = true;
else if (tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHDIST"u8)) opts.WithDist = true;
else if (tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHHASH"u8)) opts.WithHash = true;
else if (input.header.SortedSetOp == SortedSetOperation.GEOSEARCH && tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHCOORD"u8)) opts.WithCoord = true;
else if ((input.header.SortedSetOp == SortedSetOperation.GEOSEARCH && tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHDIST"u8)) ||
(input.header.SortedSetOp == SortedSetOperation.GEOSEARCHSTORE && tokenBytes.EqualsUpperCaseSpanIgnoringCase("STOREDIST"u8))) opts.WithDist = true;
else if (input.header.SortedSetOp == SortedSetOperation.GEOSEARCH && tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHHASH"u8)) opts.WithHash = true;
else if (tokenBytes.EqualsUpperCaseSpanIgnoringCase("ANY"u8)) opts.WithCountAny = true;
else
{
Expand Down
1 change: 1 addition & 0 deletions libs/server/Resp/Objects/SortedSetGeoCommands.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,7 @@ private unsafe bool GeoSearchStore<TGarnetApi>(ref TGarnetApi storageApi)
var input = new ObjectInput(new RespInputHeader
{
type = GarnetObjectType.SortedSet,
SortedSetOp = SortedSetOperation.GEOSEARCHSTORE
}, ref parseState, startIdx: 2);

var output = new SpanByteAndMemory(dcurr, (int)(dend - dcurr));
Expand Down
73 changes: 11 additions & 62 deletions libs/server/Storage/Session/ObjectStore/SortedSetGeoOps.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ public GarnetStatus GeoCommands<TObjectContext>(byte[] key, ref ObjectInput inpu
/// </summary>
/// <typeparam name="TObjectContext"></typeparam>
/// <param name="key"></param>
/// <param name="destination"></param>
/// <param name="input"></param>
/// <param name="output"></param>
/// <param name="objectContext"></param>
Expand Down Expand Up @@ -76,52 +77,10 @@ public unsafe GarnetStatus GeoSearchStore<TObjectContext>(ArgSlice key, ArgSlice

try
{
var storeDistIdx = -1;
var i = 0;
while (i < input.parseState.Count)
{
if (input.parseState.GetArgSliceByRef(i).ReadOnlySpan.EqualsUpperCaseSpanIgnoringCase(CmdStrings.STOREDIST))
{
storeDistIdx = i;
break;
}
i++;
}

SessionParseState parseState = default;
// Prepare the arguments for geo search when store distance is not provided
if (storeDistIdx == -1)
{
parseState = new SessionParseState();
parseState.Initialize(input.parseState.Count + 1);
parseState.SetArguments(0, input.parseState.Parameters);
parseState.SetArguments(input.parseState.Count, ArgSlice.FromPinnedSpan(CmdStrings.WITHHASH));
}
// When StoreDist is provided as the last argument
else if (storeDistIdx == input.parseState.Count - 1)
{
parseState = new SessionParseState();
parseState.Initialize(input.parseState.Count);
parseState.SetArguments(0, input.parseState.Parameters.Slice(0, input.parseState.Count - 1));
parseState.SetArguments(input.parseState.Count - 1, ArgSlice.FromPinnedSpan(CmdStrings.WITHDIST));
}
// When StoreDist is provided in the middle
else
{
parseState = new SessionParseState();
parseState.Initialize(input.parseState.Count);
parseState.SetArguments(0, input.parseState.Parameters.Slice(0, storeDistIdx));
parseState.SetArguments(storeDistIdx, input.parseState.Parameters.Slice(storeDistIdx + 1));
parseState.SetArguments(input.parseState.Count - 1, ArgSlice.FromPinnedSpan(CmdStrings.WITHDIST));
}

var sourceKey = key.ToArray();
var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.GEOSEARCH };
var searchInput = new ObjectInput(header, ref parseState);

SpanByteAndMemory searchOutMem = default;
var searchOut = new GarnetObjectStoreOutput { spanByteAndMemory = searchOutMem };
var status = GeoCommands(sourceKey, ref searchInput, ref searchOut, ref objectStoreLockableContext);
var status = GeoCommands(sourceKey, ref input, ref searchOut, ref objectStoreLockableContext);
searchOutMem = searchOut.spanByteAndMemory;

if (status == GarnetStatus.WRONGTYPE)
Expand All @@ -144,7 +103,7 @@ public unsafe GarnetStatus GeoSearchStore<TObjectContext>(ArgSlice key, ArgSlice
try
{
var searchOutPtr = (byte*)searchOutHandler.Pointer;
var currOutPtr = searchOutPtr;
ref var currOutPtr = ref searchOutPtr;
var endOutPtr = searchOutPtr + searchOutMem.Length;

if (RespReadUtils.TryReadErrorAsSpan(out var error, ref currOutPtr, endOutPtr))
Expand All @@ -160,35 +119,25 @@ public unsafe GarnetStatus GeoSearchStore<TObjectContext>(ArgSlice key, ArgSlice
RespReadUtils.ReadUnsignedArrayLength(out var foundItems, ref currOutPtr, endOutPtr);

// Prepare the parse state for sorted set add
var zParseState = new SessionParseState();
zParseState.Initialize(foundItems * 2);

for (int j = 0; j < foundItems; j++)
parseState.Initialize(foundItems * 2);

for (var j = 0; j < foundItems; j++)
{
RespReadUtils.ReadUnsignedArrayLength(out var innerLength, ref currOutPtr, endOutPtr);
Debug.Assert(innerLength == 2, "Should always has location and hash or distance");

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

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

var zAddOutput = new GarnetObjectStoreOutput { spanByteAndMemory = new SpanByteAndMemory(null) };
RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
Expand Down

0 comments on commit fa58e13

Please sign in to comment.