Skip to content

Commit

Permalink
GEORADIUS/GEORADIUSBYMEMBER STORE/STOREDIST
Browse files Browse the repository at this point in the history 
Fix default GEOSEARCHSTORE which apparently never worked since hash wasn't used as score.
  • Loading branch information
prvyk committed Mar 4, 2025
1 parent 4301fe9 commit 670e4ab
Show file tree
Hide file tree
Showing 6 changed files with 119 additions and 34 deletions.
47 changes: 31 additions & 16 deletions libs/server/Objects/SortedSetGeo/SortedSetGeoObjectImpl.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using Garnet.common;
using Tsavorite.core;
Expand Down Expand Up @@ -588,22 +589,36 @@ private void GeoSearch(ref ObjectInput input, ref SpanByteAndMemory output)
continue;
}

if (tokenBytes.EqualsUpperCaseSpanIgnoringCase("WITHCOORD"u8))
{
opts.withCoord = true;
continue;
}

if (tokenBytes.EqualsUpperCaseSpanIgnoringCase("ANY"u8))
{
opts.withCountAny = true;
continue;
}

// When GEORADIUS STORE support is added we'll need to account for it.
// For now we'll act as if all radius commands are readonly.
if (readOnlyCmd || byRadiusCmd)
if (!readOnlyCmd)
{
if (byRadiusCmd && tokenBytes.EqualsUpperCaseSpanIgnoringCase(CmdStrings.STORE))
{
if (byRadiusCmd)
currTokenIdx++;
continue;
}
if (tokenBytes.EqualsUpperCaseSpanIgnoringCase(CmdStrings.STOREDIST))
{
if (byRadiusCmd)
currTokenIdx++;
opts.withDist = true;
continue;
}
}
else
{
if (tokenBytes.EqualsUpperCaseSpanIgnoringCase(CmdStrings.WITHCOORD))
{
opts.withCoord = true;
continue;
}

if (tokenBytes.EqualsUpperCaseSpanIgnoringCase(CmdStrings.WITHDIST))
{
opts.withDist = true;
Expand All @@ -616,12 +631,6 @@ private void GeoSearch(ref ObjectInput input, ref SpanByteAndMemory output)
continue;
}
}
// GEOSEARCHSTORE
else if (tokenBytes.EqualsUpperCaseSpanIgnoringCase(CmdStrings.STOREDIST))
{
opts.withDist = true;
continue;
}

errorMessage = CmdStrings.RESP_SYNTAX_ERROR;
break;
Expand All @@ -646,6 +655,12 @@ private void GeoSearch(ref ObjectInput input, ref SpanByteAndMemory output)
return;
}

// On storing to ZSET, we need to use either dist or hash as score.
if (!readOnlyCmd && !opts.withDist && !opts.withHash)
{
opts.withHash = true;
}

#region act
// FROMLONLAT
bool hasLonLat = opts.origin == GeoOriginType.FromLonLat;
Expand Down Expand Up @@ -776,7 +791,7 @@ private void GeoSearch(ref ObjectInput input, ref SpanByteAndMemory output)

if (opts.withHash)
{
while (!RespWriteUtils.TryWriteInt64(item.GeoHash, ref curr, end))
while (!RespWriteUtils.TryWriteArrayItem(item.GeoHash, ref curr, end))
ObjectUtils.ReallocateOutput(ref output, ref isMemory, ref ptr, ref ptrHandle, ref curr, ref end);
}

Expand Down
2 changes: 2 additions & 0 deletions libs/server/Resp/CmdStrings.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,9 @@ static partial class CmdStrings
public static ReadOnlySpan<byte> CHANNELS => "CHANNELS"u8;
public static ReadOnlySpan<byte> NUMPAT => "NUMPAT"u8;
public static ReadOnlySpan<byte> NUMSUB => "NUMSUB"u8;
public static ReadOnlySpan<byte> STORE => "STORE"u8;
public static ReadOnlySpan<byte> STOREDIST => "STOREDIST"u8;
public static ReadOnlySpan<byte> WITHCOORD => "WITHCOORD"u8;
public static ReadOnlySpan<byte> WITHDIST => "WITHDIST"u8;
public static ReadOnlySpan<byte> WITHHASH => "WITHHASH"u8;
public static ReadOnlySpan<byte> LIB_NAME => "LIB-NAME"u8;
Expand Down
79 changes: 69 additions & 10 deletions libs/server/Resp/Objects/SortedSetGeoCommands.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ private unsafe bool GeoAdd<TGarnetApi>(ref TGarnetApi storageApi)
/// <param name="command"></param>
/// <param name="storageApi"></param>
/// <returns></returns>
private unsafe bool GeoCommands<TGarnetApi>(RespCommand command, ref TGarnetApi storageApi)
private unsafe bool GeoCommands<TGarnetApi>(RespCommand command, ref TGarnetApi storageApi, bool @readonly = true)
where TGarnetApi : IGarnetApi
{
var paramsRequiredInCommand = 0;
Expand All @@ -83,11 +83,11 @@ private unsafe bool GeoCommands<TGarnetApi>(RespCommand command, ref TGarnetApi
break;
case RespCommand.GEORADIUS:
case RespCommand.GEORADIUS_RO:
paramsRequiredInCommand = 4;
paramsRequiredInCommand = 5;
break;
case RespCommand.GEORADIUSBYMEMBER:
case RespCommand.GEORADIUSBYMEMBER_RO:
paramsRequiredInCommand = 3;
paramsRequiredInCommand = 4;
break;
}

Expand Down Expand Up @@ -136,6 +136,9 @@ private unsafe bool GeoCommands<TGarnetApi>(RespCommand command, ref TGarnetApi
throw new Exception($"Unexpected {nameof(SortedSetOperation)}: {command}");
}

if (@readonly)
opts &= ~SortedSetGeoOpts.Store;

// Prepare input
var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = op };

Expand Down Expand Up @@ -187,25 +190,81 @@ private unsafe bool GeoCommands<TGarnetApi>(RespCommand command, ref TGarnetApi
/// GEOSEARCHSTORE: Store the the members of a sorted set populated with geospatial data, which are within the borders of the area specified by a given shape.
/// </summary>
/// <typeparam name="TGarnetApi"></typeparam>
/// <param name="cmd"></param>
/// <param name="storageApi"></param>
/// <returns></returns>
private unsafe bool GeoSearchStore<TGarnetApi>(ref TGarnetApi storageApi)
private unsafe bool GeoSearchStore<TGarnetApi>(RespCommand cmd, ref TGarnetApi storageApi)
where TGarnetApi : IGarnetApi
{
// GEOSEARCHSTORE dst src FROMEMBER key BYRADIUS 0 m
if (parseState.Count < 7)
int GetDestIdx()
{
var idx = 1;
while (true)
{
if (idx >= parseState.Count - 1)
break;

var argSpan = parseState.GetArgSliceByRef(idx++).ReadOnlySpan;

if (argSpan.EqualsUpperCaseSpanIgnoringCase(CmdStrings.STORE) ||
argSpan.EqualsUpperCaseSpanIgnoringCase(CmdStrings.STOREDIST))
{
return idx;
}
}

return -1;
}

SortedSetGeoOpts opt;
int destIdx, sourceIdx;

switch (cmd)
{
return AbortWithWrongNumberOfArguments(nameof(RespCommand.GEOSEARCHSTORE));
case RespCommand.GEORADIUS:
// GERADIUS src lon lat 0 km
if (parseState.Count < 5)
{
return AbortWithWrongNumberOfArguments(nameof(RespCommand.GEORADIUS));
}
opt = SortedSetGeoOpts.Store | SortedSetGeoOpts.ByRadius;
sourceIdx = 0;
destIdx = GetDestIdx();
break;
case RespCommand.GEORADIUSBYMEMBER:
// GERADIUSBYMEMBER src member 0 km
if (parseState.Count < 4)
{
return AbortWithWrongNumberOfArguments(nameof(RespCommand.GEORADIUSBYMEMBER));
}
opt = SortedSetGeoOpts.Store | SortedSetGeoOpts.ByRadius | SortedSetGeoOpts.ByMember;
sourceIdx = 0;
destIdx = GetDestIdx();
break;
case RespCommand.GEOSEARCHSTORE:
default:
// GEOSEARCHSTORE dst src FROMEMBER key BYRADIUS 0 m
if (parseState.Count < 7)
{
return AbortWithWrongNumberOfArguments(nameof(RespCommand.GEOSEARCHSTORE));
}
opt = SortedSetGeoOpts.Store;
destIdx = 0;
sourceIdx = 1;
break;
}

var destinationKey = parseState.GetArgSliceByRef(0);
var sourceKey = parseState.GetArgSliceByRef(1);
if (destIdx == -1)
return GeoCommands(cmd, ref storageApi, true);

var destinationKey = parseState.GetArgSliceByRef(destIdx);
var sourceKey = parseState.GetArgSliceByRef(sourceIdx);

var input = new ObjectInput(new RespInputHeader
{
type = GarnetObjectType.SortedSet,
SortedSetOp = SortedSetOperation.GEOSEARCH,
}, ref parseState, startIdx: 2, arg2: (int)SortedSetGeoOpts.Store);
}, ref parseState, startIdx: sourceIdx + 1, arg2: (int)opt);

var output = new SpanByteAndMemory(dcurr, (int)(dend - dcurr));
var status = storageApi.GeoSearchStore(sourceKey, destinationKey, ref input, ref output);
Expand Down
8 changes: 4 additions & 4 deletions libs/server/Resp/RespServerSession.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -726,12 +726,12 @@ private bool ProcessArrayCommands<TGarnetApi>(RespCommand cmd, ref TGarnetApi st
RespCommand.GEOHASH => GeoCommands(cmd, ref storageApi),
RespCommand.GEODIST => GeoCommands(cmd, ref storageApi),
RespCommand.GEOPOS => GeoCommands(cmd, ref storageApi),
RespCommand.GEORADIUS_RO or
RespCommand.GEORADIUS => GeoSearchStore(cmd, ref storageApi),
RespCommand.GEORADIUS_RO => GeoCommands(cmd, ref storageApi),
RespCommand.GEORADIUSBYMEMBER => GeoSearchStore(cmd, ref storageApi),
RespCommand.GEORADIUSBYMEMBER_RO => GeoCommands(cmd, ref storageApi),
RespCommand.GEORADIUS or
RespCommand.GEORADIUSBYMEMBER => GeoCommands(cmd, ref storageApi),
RespCommand.GEOSEARCH => GeoCommands(cmd, ref storageApi),
RespCommand.GEOSEARCHSTORE => GeoSearchStore(ref storageApi),
RespCommand.GEOSEARCHSTORE => GeoSearchStore(cmd, ref storageApi),
//HLL Commands
RespCommand.PFADD => HyperLogLogAdd(ref storageApi),
RespCommand.PFMERGE => HyperLogLogMerge(ref storageApi),
Expand Down
9 changes: 9 additions & 0 deletions test/Garnet.test/RespSortedSetGeoTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -400,6 +400,9 @@ public void CanUseGeoSearchStore()
Assert.That(actualValues[1].Score, Is.EqualTo(198.424300439725).Within(1.0 / Math.Pow(10, 6)));
ClassicAssert.AreEqual("New York", (string)actualValues[2].Element);
Assert.That(actualValues[2].Score, Is.EqualTo(327.676458633557).Within(1.0 / Math.Pow(10, 6)));

actualCount = db.GeoSearchAndStore(key, destinationKey, new RedisValue("Washington"), new GeoSearchBox(800, 800, GeoUnit.Kilometers), storeDistances: true);
ClassicAssert.AreEqual(3, actualCount);
}

[Test]
Expand All @@ -418,6 +421,12 @@ public void CanUseGeoSearchStoreWithDeleteKeyWhenSourceNotFound()

var actualValues = db.SortedSetRangeByScoreWithScores(destinationKey);
ClassicAssert.AreEqual(0, actualValues.Length);

actualCount = db.GeoSearchAndStore(key, destinationKey, new RedisValue("Washington"), new GeoSearchBox(800, 800, GeoUnit.Kilometers));
ClassicAssert.AreEqual(0, actualCount);

actualValues = db.SortedSetRangeByScoreWithScores(destinationKey);
ClassicAssert.AreEqual(0, actualValues.Length);
}

//end region of SE tests
Expand Down
8 changes: 4 additions & 4 deletions website/docs/commands/api-compatibility.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,12 +155,12 @@ Note that this list is subject to change as we continue to expand our API comman
| | [GEODIST](data-structures.md#geodist) | ➕ | |
| | [GEOHASH](data-structures.md#geohash) | ➕ | |
| | [GEOPOS](data-structures.md#geopos) | ➕ | |
| | [GEORADIUS](data-structures.md#georadius) | ➕ | (Deprecated) Partially Implemented |
| | [GEORADIUS](data-structures.md#georadius) | ➕ | (Deprecated) |
| | [GEORADIUS_RO](data-structures.md#georadius_ro) | ➕ | (Deprecated) |
| | [GEORADIUSBYMEMBER](data-structures.md#georadiusbymember) | ➕ | (Deprecated) Partially Implemented |
| | [GEORADIUSBYMEMBER](data-structures.md#georadiusbymember) | ➕ | (Deprecated) |
| | [GEORADIUSBYMEMBER_RO](data-structures.md#georadiusbymember_ro) | ➕ | (Deprecated) |
| | [GEOSEARCH](data-structures.md#geosearch) | ➕ | Partially Implemented |
| | [GEOSEARCHSTORE](data-structures.md#geosearchstore) | ➕ | Partially Implemented |
| | [GEOSEARCH](data-structures.md#geosearch) | ➕ | |
| | [GEOSEARCHSTORE](data-structures.md#geosearchstore) | ➕ | |
| <span id="hash">**HASH**</span> | [HDEL](data-structures.md#hdel) | ➕ | |
| | [HEXISTS](data-structures.md#hexists) | ➕ | |
| | [HEXPIRE](data-structures.md#hexpire) | ➕ | |
Expand Down

0 comments on commit 670e4ab

Please sign in to comment.