diff --git a/go.mod b/go.mod index dcf33b828..2f7781b61 100644 --- a/go.mod +++ b/go.mod @@ -43,7 +43,6 @@ require ( github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.3 github.com/mattn/go-sqlite3 v1.14.24 - github.com/mmcloughlin/geohash v0.10.0 github.com/ohler55/ojg v1.25.0 github.com/rs/xid v1.6.0 github.com/rs/zerolog v1.33.0 diff --git a/go.sum b/go.sum index 51567b272..ebfb5659b 100644 --- a/go.sum +++ b/go.sum @@ -74,8 +74,6 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-sqlite3 v1.14.24 h1:tpSp2G2KyMnnQu99ngJ47EIkWVmliIizyZBfPrBWDRM= github.com/mattn/go-sqlite3 v1.14.24/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= -github.com/mmcloughlin/geohash v0.10.0 h1:9w1HchfDfdeLc+jFEf/04D27KP7E2QmpDu52wPbJWRE= -github.com/mmcloughlin/geohash v0.10.0/go.mod h1:oNZxQo5yWJh0eMQEP/8hwQuVx9Z9tjwFUqcTB1SmG0c= github.com/ohler55/ojg v1.25.0 h1:sDwc4u4zex65Uz5Nm7O1QwDKTT+YRcpeZQTy1pffRkw= github.com/ohler55/ojg v1.25.0/go.mod h1:gQhDVpQLqrmnd2eqGAvJtn+NfKoYJbe/A4Sj3/Vro4o= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= diff --git a/internal/clientio/resp.go b/internal/clientio/resp.go index fdfc1dccb..139459450 100644 --- a/internal/clientio/resp.go +++ b/internal/clientio/resp.go @@ -236,6 +236,16 @@ func Encode(value interface{}, isSimple bool) []byte { buf.Write(encodeString(b)) // Encode each string and write to the buffer. } return []byte(fmt.Sprintf("*%d\r\n%s", len(v), buf.Bytes())) // Return the encoded response. + case [][]interface{}: + var b []byte + buf := bytes.NewBuffer(b) + + buf.WriteString(fmt.Sprintf("*%d\r\n", len(v))) + + for _, list := range v { + buf.Write(Encode(list, false)) + } + return buf.Bytes() // Handle slices of custom objects (Obj). case []*object.Obj: @@ -255,6 +265,15 @@ func Encode(value interface{}, isSimple bool) []byte { } return []byte(fmt.Sprintf("*%d\r\n%s", len(v), buf.Bytes())) // Return the encoded response. + // Handle slices of int64. + case []float64: + var b []byte + buf := bytes.NewBuffer(b) // Create a buffer for accumulating encoded values. + for _, b := range value.([]float64) { + buf.Write(Encode(b, false)) // Encode each int64 and write to the buffer. + } + return []byte(fmt.Sprintf("*%d\r\n%s", len(v), buf.Bytes())) // Return the encoded response. + // Handle slices of int64. case []int64: var b []byte diff --git a/internal/errors/migrated_errors.go b/internal/errors/migrated_errors.go index ea81b0c4e..c299ed8ea 100644 --- a/internal/errors/migrated_errors.go +++ b/internal/errors/migrated_errors.go @@ -34,6 +34,8 @@ var ( ErrInvalidFingerprint = errors.New("invalid fingerprint") ErrKeyDoesNotExist = errors.New("ERR could not perform this operation on a key that doesn't exist") ErrKeyExists = errors.New("ERR key exists") + ErrInvalidFloat = errors.New("ERR value is not a valid float") + ErrUnsupportedUnit = errors.New("ERR unsupported unit provided. please use m, km, ft, mi") // Error generation functions for specific error messages with dynamic parameters. ErrWrongArgumentCount = func(command string) error { diff --git a/internal/eval/commands.go b/internal/eval/commands.go index 787bca20c..724a4fcb8 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -1301,6 +1301,14 @@ var ( NewEval: evalGEODIST, KeySpecs: KeySpecs{BeginIndex: 1}, } + geoRadiusByMemberCmdMeta = DiceCmdMeta{ + Name: "GEORADIUSBYMEMBER", + Info: `Returns all members within a radius of a given member from the geospatial index.`, + Arity: -4, + IsMigrated: true, + NewEval: evalGEORADIUSBYMEMBER, + KeySpecs: KeySpecs{BeginIndex: 1}, + } jsonstrappendCmdMeta = DiceCmdMeta{ Name: "JSON.STRAPPEND", Info: `JSON.STRAPPEND key [path] value @@ -1444,6 +1452,7 @@ func init() { DiceCmds["FLUSHDB"] = flushdbCmdMeta DiceCmds["GEOADD"] = geoAddCmdMeta DiceCmds["GEODIST"] = geoDistCmdMeta + DiceCmds["GEORADIUSBYMEMBER"] = geoRadiusByMemberCmdMeta DiceCmds["GET"] = getCmdMeta DiceCmds["GETBIT"] = getBitCmdMeta DiceCmds["GETDEL"] = getDelCmdMeta diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index f1a629a6e..73be73400 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "math" + "math/rand" "reflect" "strconv" "strings" @@ -132,6 +133,7 @@ func TestEval(t *testing.T) { testEvalBitFieldRO(t, store) testEvalGEOADD(t, store) testEvalGEODIST(t, store) + testEvalGEORADIUSBYMEMBER(t, store) testEvalSINTER(t, store) testEvalJSONSTRAPPEND(t, store) testEvalINCR(t, store) @@ -8295,7 +8297,7 @@ func testEvalGEODIST(t *testing.T, store *dstore.Store) { }, input: []string{"points", "Palermo", "Catania"}, migratedOutput: EvalResponse{ - Result: float64(166274.1440), + Result: float64(166274.1516), Error: nil, }, }, @@ -8306,7 +8308,7 @@ func testEvalGEODIST(t *testing.T, store *dstore.Store) { }, input: []string{"points", "Palermo", "Catania", "km"}, migratedOutput: EvalResponse{ - Result: float64(166.2741), + Result: float64(166.2742), Error: nil, }, }, @@ -9222,3 +9224,301 @@ func testEvalLRANGE(t *testing.T, store *dstore.Store) { } runMigratedEvalTests(t, tests, evalLRANGE, store) } + +func testEvalGEORADIUSBYMEMBER(t *testing.T, store *dstore.Store) { + tests := map[string]evalTestCase{ + "GEORADIUSBYMEMBER wrong number of arguments": { + input: []string{"nyc", "wtc one", "7"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEORADIUSBYMEMBER"), + }, + }, + "GEORADIUSBYMEMBER non-numeric radius": { + input: []string{"nyc", "wtc one", "invalid", "km"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("need numeric radius"), + }, + }, + "GEORADIUSBYMEMBER non-existing key": { + input: []string{"nonexistent", "wtc one", "7", "km"}, + migratedOutput: EvalResponse{ + Result: clientio.EmptyArray, + Error: nil, + }, + }, + "GEORADIUSBYMEMBER wrong type operation": { + setup: func() { + store.Put("wrongtype", store.NewObj("string_value", -1, object.ObjTypeString)) + }, + input: []string{"wrongtype", "wtc one", "7", "km"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + }, + }, + "GEORADIUSBYMEMBER non-existing member": { + setup: func() { + evalGEOADD([]string{"nyc", "-73.9798091", "40.7598464", "wtc one"}, store) + }, + input: []string{"nyc", "nonexistent", "7", "km"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("could not decode requested zset member"), + }, + }, + "GEORADIUSBYMEMBER unsupported unit": { + setup: func() { + evalGEOADD([]string{"nyc", "-73.9798091", "40.7598464", "wtc one"}, store) + }, + input: []string{"nyc", "wtc one", "7", "invalid"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrUnsupportedUnit, + }, + }, + "GEORADIUSBYMEMBER simple radius search": { + setup: func() { + evalGEOADD([]string{"nyc", + "-73.9798091", "40.7598464", "wtc one", + "-73.981", "40.768", "union square", + "-73.973", "40.764", "central park n/q/r", + "-73.990", "40.750", "4545", + "-73.953", "40.748", "lic market", + }, store) + }, + input: []string{"nyc", "wtc one", "7", "km"}, + migratedOutput: EvalResponse{ + Result: []string{"wtc one", "union square", "central park n/q/r", "4545", "lic market"}, + Error: nil, + }, + }, + "GEORADIUSBYMEMBER oblique direction search close points": { + setup: func() { + evalGEOADD([]string{"k1", + "-0.15307903289794921875", "85", "n1", + "0.3515625", "85.00019260486917005437", "n2", + }, store) + }, + input: []string{"k1", "n1", "4891.94", "m"}, + migratedOutput: EvalResponse{ + Result: []string{"n1", "n2"}, + Error: nil, + }, + }, + "GEORADIUSBYMEMBER oblique direction search distant points": { + setup: func() { + evalGEOADD([]string{"k1", + "-4.95211958885192871094", "85", "n3", + "11.25", "85.0511", "n4", + }, store) + }, + input: []string{"k1", "n3", "156544", "m"}, + migratedOutput: EvalResponse{ + Result: []string{"n3", "n4"}, + Error: nil, + }, + }, + "GEORADIUSBYMEMBER crossing poles": { + setup: func() { + evalGEOADD([]string{"k1", + "45", "65", "n1", + "-135", "85.05", "n2", + }, store) + }, + input: []string{"k1", "n1", "5009431", "m"}, + migratedOutput: EvalResponse{ + Result: []string{"n1", "n2"}, + Error: nil, + }, + }, + "GEORADIUSBYMEMBER with coordinates option": { + setup: func() { + evalGEOADD([]string{"nyc", "-73.9798091", "40.7598464", "wtc one"}, store) + }, + input: []string{"nyc", "wtc one", "7", "km", "WITHCOORD"}, + migratedOutput: EvalResponse{ + Result: [][]interface{}{ + {"wtc one", []float64{40.759845946389994, -73.97980660200119}}, + }, + Error: nil, + }, + }, + "GEORADIUSBYMEMBER with distance option": { + setup: func() { + evalGEOADD([]string{"nyc", "-73.9798091", "40.7598464", "wtc one"}, store) + }, + input: []string{"nyc", "wtc one", "7", "km", "WITHDIST"}, + migratedOutput: EvalResponse{ + Result: [][]interface{}{ + {"wtc one", 0.0}, + }, + Error: nil, + }, + }, + } + + runMigratedEvalTests(t, tests, evalGEORADIUSBYMEMBER, store) +} + +func generateMembers(count int) []struct { + name string + latitude float64 + longitude float64 +} { + locations := make([]struct { + name string + latitude float64 + longitude float64 + }, count) + + // Generate locations around San Francisco area + for i := range locations { + locations[i] = struct { + name string + latitude float64 + longitude float64 + }{ + name: "loc" + strconv.Itoa(i), + latitude: 37.7749 + (rand.Float64()*2 - 1), // ±1 degree from SF + longitude: -122.4194 + (rand.Float64()*2 - 1), + } + } + return locations +} + +// setupTestStore creates a store with test data +func setupTestStore(locations []struct { + name string + latitude float64 + longitude float64 +}) *dstore.Store { + store := dstore.NewStore(nil, nil, nil) + + for _, loc := range locations { + evalGEOADD([]string{ + "locations", + "NX", + fmt.Sprintf("%f", loc.longitude), + fmt.Sprintf("%f", loc.latitude), + loc.name, + }, store) + } + + return store +} + +func BenchmarkGEORADIUSBYMEMBER(b *testing.B) { + scenarios := []struct { + name string + locations int + }{ + {"Small", 100}, + {"Medium", 1000}, + {"Large", 10000}, + } + + for _, sc := range scenarios { + b.Run(sc.name, func(b *testing.B) { + locations := generateMembers(sc.locations) + store := setupTestStore(locations) + + middlePoint := locations[len(locations)/2] + + args := []string{ + "locations", + middlePoint.name, + "5", + "km", + "WITHCOORD", + "WITHDIST", + "COUNT", "50", + "ASC", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + response := evalGEORADIUSBYMEMBER(args, store) + if response.Error != nil { + b.Fatal(response.Error) + } + } + }) + } +} + +// BenchmarkGEORADIUSBYMEMBER_DifferentRadii tests performance with different search radii +func BenchmarkGEORADIUSBYMEMBER_DifferentRadii(b *testing.B) { + cases := []struct { + name string + radius string + }{ + {"SmallRadius", "1"}, + {"MediumRadius", "10"}, + {"LargeRadius", "50"}, + } + + locations := generateMembers(1000) + store := setupTestStore(locations) + middlePoint := locations[len(locations)/2] + + for _, tc := range cases { + b.Run(tc.name, func(b *testing.B) { + args := []string{ + "locations", + middlePoint.name, + tc.radius, + "km", + "WITHCOORD", + "WITHDIST", + "ASC", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + response := evalGEORADIUSBYMEMBER(args, store) + if response.Error != nil { + b.Fatal(response.Error) + } + } + }) + } +} + +// BenchmarkGEORADIUSBYMEMBER_DifferentOptions tests performance with different output options +func BenchmarkGEORADIUSBYMEMBER_DifferentOptions(b *testing.B) { + locations := generateMembers(1000) + store := setupTestStore(locations) + middlePoint := locations[len(locations)/2] + + cases := []struct { + name string + options []string + }{ + {"NoOptions", []string{}}, + {"WithCoord", []string{"WITHCOORD"}}, + {"WithDist", []string{"WITHDIST"}}, + {"WithHash", []string{"WITHHASH"}}, + {"AllOptions", []string{"WITHCOORD", "WITHDIST", "WITHHASH"}}, + } + + for _, tc := range cases { + b.Run(tc.name, func(b *testing.B) { + args := append([]string{ + "locations", + middlePoint.name, + "5", + "km", + }, tc.options...) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + response := evalGEORADIUSBYMEMBER(args, store) + if response.Error != nil { + b.Fatal(response.Error) + } + } + }) + } +} diff --git a/internal/eval/geo/geo.go b/internal/eval/geo/geo.go index 6db40eaf8..0cfd2e776 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -3,30 +3,47 @@ package geo import ( "math" - "github.com/dicedb/dice/internal/errors" - "github.com/mmcloughlin/geohash" + diceerrors "github.com/dicedb/dice/internal/errors" ) // Earth's radius in meters const earthRadius float64 = 6372797.560856 -// Bit precision for geohash - picked up to match redis -const bitPrecision = 52 +// Bit precision steps for geohash - picked up to match redis +const maxSteps = 26 +const mercatorMax = 20037726.37 + +const ( + /* These are constraints from EPSG:900913 / EPSG:3785 / OSGEO:41001 */ + /* We can't geocode at the north/south pole. */ + globalMinLat = -85.05112878 + globalMaxLat = 85.05112878 + globalMinLon = -180.0 + globalMaxLon = 180.0 +) + +type Unit string + +const ( + Meters Unit = "m" + Kilometers Unit = "km" + Miles Unit = "mi" + Feet Unit = "ft" +) + +// DegToRad converts degrees to radians. func DegToRad(deg float64) float64 { return math.Pi * deg / 180.0 } +// RadToDeg converts radians to degrees. func RadToDeg(rad float64) float64 { return 180.0 * rad / math.Pi } -func GetDistance( - lon1, - lat1, - lon2, - lat2 float64, -) float64 { +// GetDistance calculates the distance between two geographical points specified by their longitude and latitude. +func GetDistance(lon1, lat1, lon2, lat2 float64) float64 { lon1r := DegToRad(lon1) lon2r := DegToRad(lon2) v := math.Sin((lon2r - lon1r) / 2) @@ -44,43 +61,386 @@ func GetDistance( return 2.0 * earthRadius * math.Asin(math.Sqrt(a)) } +// GetLatDistance calculates the distance between two latitudes. func GetLatDistance(lat1, lat2 float64) float64 { return earthRadius * math.Abs(DegToRad(lat2)-DegToRad(lat1)) } -// EncodeHash returns a geo hash for a given coordinate, and returns it in float64 so it can be used as score in a zset -func EncodeHash( - latitude, - longitude float64, -) float64 { - h := geohash.EncodeIntWithPrecision(latitude, longitude, bitPrecision) +// EncodeHash returns a geo hash for a given coordinate, and returns it in float64 so it can be used as score in a zset. +func EncodeHash(latitude, longitude float64) float64 { + h := encodeHash(longitude, latitude, maxSteps) + h = align52Bits(h, maxSteps) return float64(h) } -// DecodeHash returns the latitude and longitude from a geo hash -// The hash should be a float64, as it is used as score in a zset +// encodeHash encodes the latitude and longitude into a geohash with the specified number of steps. +func encodeHash(longitude, latitude float64, steps uint8) uint64 { + latOffset := (latitude - globalMinLat) / (globalMaxLat - globalMinLat) + longOffset := (longitude - globalMinLon) / (globalMaxLon - globalMinLon) + + latOffset *= float64(uint64(1) << steps) + longOffset *= float64(uint64(1) << steps) + return interleave64(uint32(latOffset), uint32(longOffset)) +} + +// DecodeHash returns the latitude and longitude from a geo hash. +// The hash should be a float64, as it is used as score in a sorted set. func DecodeHash(hash float64) (lat, lon float64) { - lat, lon = geohash.DecodeIntWithPrecision(uint64(hash), bitPrecision) + return decodeHash(uint64(hash), maxSteps) +} + +// decodeHash decodes the geohash into latitude and longitude with the specified number of steps. +func decodeHash(hash uint64, steps uint8) (lat, lon float64) { + hashSep := deinterleave64(hash) + + latScale := globalMaxLat - globalMinLat + longScale := globalMaxLon - globalMinLon + + ilato := uint32(hashSep) // lat part + ilono := uint32(hashSep >> 32) // lon part + + // divide by 2**step. + // Then, for 0-1 coordinate, multiply times scale and add + // to the min to get the absolute coordinate. + minLat := globalMinLat + (float64(ilato)*1.0/float64(uint64(1)< globalMaxLon { + lon = globalMaxLon + } + if lon < globalMinLon { + lon = globalMinLon + } + + lat = (minLat + maxLat) / 2 + if lat > globalMaxLat { + lat = globalMaxLat + } + if lat < globalMinLat { + lat = globalMinLat + } return lat, lon } -// ConvertDistance converts a distance from meters to the desired unit -func ConvertDistance( - distance float64, - unit string, -) (converted float64, err []byte) { - switch unit { - case "m": +// ConvertDistance converts a distance from meters to the desired unit. +func ConvertDistance(distance float64, unit string) (float64, error) { + switch Unit(unit) { + case Meters: return distance, nil - case "km": + case Kilometers: return distance / 1000, nil - case "mi": + case Miles: return distance / 1609.34, nil - case "ft": + case Feet: return distance / 0.3048, nil default: - return 0, errors.NewErrWithMessage("ERR unsupported unit provided. please use m, km, ft, mi") + return 0, diceerrors.ErrUnsupportedUnit + } +} + +// ToMeters converts a distance and its unit to meters. +func ToMeters(distance float64, unit string) (float64, bool) { + switch Unit(unit) { + case Meters: + return distance, true + case Kilometers: + return distance * 1000, true + case Miles: + return distance * 1609.34, true + case Feet: + return distance * 0.3048, true + default: + return 0, false + } +} + +// geohashEstimateStepsByRadius estimates the number of steps required to cover a radius at a given latitude. +func geohashEstimateStepsByRadius(radius, lat float64) uint8 { + if radius == 0 { + return 26 + } + + step := 1 + for radius < mercatorMax { + radius *= 2 + step++ + } + step -= 2 // Make sure range is included in most of the base cases. + + /* Note from the redis implementation: + Wider range towards the poles... Note: it is possible to do better + than this approximation by computing the distance between meridians + at this latitude, but this does the trick for now. */ + if lat > 66 || lat < -66 { + step-- + if lat > 80 || lat < -80 { + step-- + } + } + + if step < 1 { + step = 1 + } + if step > 26 { + step = 26 + } + + return uint8(step) +} + +// boundingBox returns the bounding box for a given latitude, longitude and radius. +func boundingBox(lat, lon, radius float64) (minLon, minLat, maxLon, maxLat float64) { + latDelta := RadToDeg(radius / earthRadius) + lonDeltaTop := RadToDeg(radius / earthRadius / math.Cos(DegToRad(lat+latDelta))) + lonDeltaBottom := RadToDeg(radius / earthRadius / math.Cos(DegToRad(lat-latDelta))) + + isSouthernHemisphere := false + if lat < 0 { + isSouthernHemisphere = true + } + + minLon = lon - lonDeltaTop + if isSouthernHemisphere { + minLon = lon - lonDeltaBottom + } + + maxLon = lon + lonDeltaTop + if isSouthernHemisphere { + maxLon = lon + lonDeltaBottom + } + + minLat = lat - latDelta + maxLat = lat + latDelta + + return minLon, minLat, maxLon, maxLat +} + +// Area returns the geohashes of the area covered by a circle with a given radius. It returns the center hash +// and the 8 surrounding hashes. The second return value is the number of steps used to cover the area. +func Area(centerHash, radius float64) (result [9]uint64, steps uint8) { + centerLat, centerLon := decodeHash(uint64(centerHash), maxSteps) + minLon, minLat, maxLon, maxLat := boundingBox(centerLat, centerLon, radius) + steps = geohashEstimateStepsByRadius(radius, centerLat) + centerRadiusHash := encodeHash(centerLon, centerLat, steps) + + neighbors := geohashNeighbors(centerRadiusHash, steps) + area := areaBySteps(centerRadiusHash, steps) + + /* Check if the step is enough at the limits of the covered area. + * Sometimes when the search area is near an edge of the + * area, the estimated step is not small enough, since one of the + * north / south / west / east square is too near to the search area + * to cover everything. */ + north := areaBySteps(neighbors[0], steps) + south := areaBySteps(neighbors[1], steps) + east := areaBySteps(neighbors[2], steps) + west := areaBySteps(neighbors[3], steps) + + decreaseStep := false + if north.Lat.Max < maxLat || south.Lat.Min > minLat || east.Lon.Max < maxLon || west.Lon.Min > minLon { + decreaseStep = true + } + + if steps > 1 && decreaseStep { + steps-- + centerRadiusHash = encodeHash(centerLon, centerLat, steps) + neighbors = geohashNeighbors(centerRadiusHash, steps) + area = areaBySteps(centerRadiusHash, steps) + } + + // exclude useless areas + if steps >= 2 { + if area.Lat.Min < minLat { + neighbors[6] = 0 // south east + neighbors[1] = 0 // south + neighbors[7] = 0 // south west + } + + if area.Lat.Max > maxLat { + neighbors[0] = 0 // north + neighbors[4] = 0 // north east + neighbors[5] = 0 // north west + } + + if area.Lon.Min < minLon { + neighbors[7] = 0 // south west + neighbors[3] = 0 // west + neighbors[5] = 0 // north west + } + + if area.Lon.Max > maxLon { + neighbors[4] = 0 // north east + neighbors[2] = 0 // east + neighbors[6] = 0 // south east + } + } + + result[0] = centerRadiusHash + for i := 0; i < len(neighbors); i++ { + result[i+1] = neighbors[i] + } + + return result, steps +} + +// HashMinMax returns the min and max hashes for a given hash and steps. This can be used to get the range of hashes +// that a given hash and a radius (steps) will cover. +func HashMinMax(hash uint64, steps uint8) (minHash, maxHash uint64) { + minHash = align52Bits(hash, steps) + hash++ + maxHash = align52Bits(hash, steps) + return minHash, maxHash +} + +// align52Bits aligns the hash to 52 bits. +func align52Bits(hash uint64, steps uint8) uint64 { + hash <<= (52 - steps*2) + return hash +} + +type hashRange struct { + Min float64 + Max float64 +} + +type hashArea struct { + Lat hashRange + Lon hashRange +} + +// deinterleave64 deinterleaves a 64-bit integer. +func deinterleave64(interleaved uint64) uint64 { + x := interleaved & 0x5555555555555555 + y := (interleaved >> 1) & 0x5555555555555555 + + x = (x | (x >> 1)) & 0x3333333333333333 + y = (y | (y >> 1)) & 0x3333333333333333 + + x = (x | (x >> 2)) & 0x0f0f0f0f0f0f0f0f + y = (y | (y >> 2)) & 0x0f0f0f0f0f0f0f0f + + x = (x | (x >> 4)) & 0x00ff00ff00ff00ff + y = (y | (y >> 4)) & 0x00ff00ff00ff00ff + + x = (x | (x >> 8)) & 0x0000ffff0000ffff + y = (y | (y >> 8)) & 0x0000ffff0000ffff + + x = (x | (x >> 16)) & 0x00000000ffffffff + y = (y | (y >> 16)) & 0x00000000ffffffff + + return (y << 32) | x +} + +// interleave64 interleaves two 32-bit integers into a 64-bit integer. +func interleave64(xlo, ylo uint32) uint64 { + B := []uint64{ + 0x5555555555555555, + 0x3333333333333333, + 0x0F0F0F0F0F0F0F0F, + 0x00FF00FF00FF00FF, + 0x0000FFFF0000FFFF, } + S := []uint{1, 2, 4, 8, 16} + + x := uint64(xlo) + y := uint64(ylo) + + x = (x | (x << S[4])) & B[4] + y = (y | (y << S[4])) & B[4] + + x = (x | (x << S[3])) & B[3] + y = (y | (y << S[3])) & B[3] + + x = (x | (x << S[2])) & B[2] + y = (y | (y << S[2])) & B[2] + + x = (x | (x << S[1])) & B[1] + y = (y | (y << S[1])) & B[1] + + x = (x | (x << S[0])) & B[0] + y = (y | (y << S[0])) & B[0] + + return x | (y << 1) +} + +// areaBySteps calculates the area covered by a hash at a given number of steps. +func areaBySteps(hash uint64, steps uint8) *hashArea { + hashSep := deinterleave64(hash) + + latScale := globalMaxLat - globalMinLat + longScale := globalMaxLon - globalMinLon + + ilato := uint32(hashSep) // lat part + ilono := uint32(hashSep >> 32) // lon part + + // divide by 2**step. + // Then, for 0-1 coordinate, multiply times scale and add + // to the min to get the absolute coordinate. + area := &hashArea{} + area.Lat.Min = globalMinLat + (float64(ilato)/float64(uint64(1)<> (64 - steps*2) + + if d > 0 { + x += uint64(zz + 1) + } else { + x |= uint64(zz) + x -= uint64(zz + 1) + } + + x &= (0xaaaaaaaaaaaaaaaa >> (64 - steps*2)) + return x | y +} + +// geohashMoveY moves the geohash in the y direction. +func geohashMoveY(hash uint64, steps uint8, d int8) uint64 { + x := hash & 0xaaaaaaaaaaaaaaaa + y := hash & 0x5555555555555555 + + zz := uint64(0xaaaaaaaaaaaaaaaa) >> (64 - steps*2) + + if d > 0 { + y += (zz + 1) + } else { + y |= zz + y -= (zz + 1) + } + + y &= (0x5555555555555555 >> (64 - steps*2)) + return x | y +} + +// geohashNeighbors returns the geohash neighbors of a given hash with a given number of steps. +func geohashNeighbors(hash uint64, steps uint8) [8]uint64 { + neighbors := [8]uint64{} + + neighbors[0] = geohashMoveY(hash, steps, 1) // North + neighbors[1] = geohashMoveY(hash, steps, -1) // South + neighbors[2] = geohashMoveX(hash, steps, 1) // East + neighbors[3] = geohashMoveX(hash, steps, -1) // West + neighbors[4] = geohashMoveX(geohashMoveY(hash, steps, 1), steps, 1) // North-East + neighbors[5] = geohashMoveX(geohashMoveY(hash, steps, 1), steps, -1) // North-West + neighbors[6] = geohashMoveX(geohashMoveY(hash, steps, -1), steps, 1) // South-East + neighbors[7] = geohashMoveX(geohashMoveY(hash, steps, -1), steps, -1) // South-West + + return neighbors } diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index 5b4435c3f..cde8e575f 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -106,7 +106,7 @@ func (ss *Set) Remove(member string) bool { return true } -// GetRange returns a slice of members with scores between min and max, inclusive. +// GetRange returns a slice of members with indices between min and max, inclusive. // it returns the members in ascending order if reverse is false, and descending order if reverse is true. // If withScores is true, the members will be returned with their scores. func (ss *Set) GetRange( @@ -309,3 +309,29 @@ func DeserializeSortedSet(buf *bytes.Reader) (*Set, error) { return ss, nil } + +// GetScoreRange returns a slice of members with scores between min and max, inclusive. +// If withScores is true, the members will be returned with their scores. +func (ss *Set) GetMemberScoresInRange(minScore, maxScore float64, count, maxCount int) (members []string, scores []float64) { + iterFunc := func(item btree.Item) bool { + ssi := item.(*Item) + if ssi.Score < minScore { + return true + } + if ssi.Score >= maxScore { + return false + } + members = append(members, ssi.Member) + scores = append(scores, ssi.Score) + count++ + + if maxCount > 0 && count >= maxCount { + return false + } + + return true + } + + ss.tree.Ascend(iterFunc) + return members, scores +} diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index c80ad9d63..25dd6effd 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6207,12 +6207,12 @@ func evalGEODIST(args []string, store *dstore.Store) *EvalResponse { distance := geo.GetDistance(lon1, lat1, lon2, lat2) - result, err := geo.ConvertDistance(distance, unit) + result, conversionErr := geo.ConvertDistance(distance, unit) - if err != nil { + if conversionErr != nil { return &EvalResponse{ Result: nil, - Error: diceerrors.ErrWrongTypeOperation, + Error: conversionErr, } } @@ -6772,3 +6772,273 @@ func evalCommandDocs(args []string) *EvalResponse { return makeEvalResult(result) } + +type geoRadiusOpts struct { + WithCoord bool + WithDist bool + WithHash bool + Count int // 0 means no count specified + CountAny bool // true if ANY was specified with COUNT + IsSorted bool // By default return items are not sorted + Ascending bool // If IsSorted is true, return items nearest to farthest relative to the center (ascending) or farthest to nearest relative to the center (descending) + Store string // If both StoreDist and Store are specified, last argument takes precedence + StoreDist bool +} + +func parseGeoRadiusOpts(args []string) (*geoRadiusOpts, error) { + opts := &geoRadiusOpts{} + + for i := 0; i < len(args); i++ { + option := strings.ToUpper(args[i]) + + switch option { + case "WITHCOORD": + opts.WithCoord = true + case "WITHDIST": + opts.WithDist = true + case "WITHHASH": + opts.WithHash = true + case "ASC": + opts.IsSorted = true + opts.Ascending = true + case "DESC": + opts.IsSorted = true + opts.Ascending = false + case "COUNT": + if i+1 < len(args) { + count, err := strconv.Atoi(args[i+1]) + if err != nil { + return nil, diceerrors.ErrIntegerOutOfRange + } + opts.Count = count + if i+2 < len(args) && strings.EqualFold(args[i+2], "ANY") { + opts.CountAny = true + i++ + } + i++ + } else { + return nil, diceerrors.ErrSyntax + } + case "ANY": + return nil, diceerrors.ErrGeneral("the ANY argument requires COUNT argument") + case "STORE": + if opts.WithCoord || opts.WithDist || opts.WithHash { + return nil, diceerrors.ErrGeneral("STORE option in GEORADIUS is not compatible with WITHDIST, WITHHASH and WITHCOORD options") + } + if i+1 < len(args) { + opts.Store = args[i+1] + opts.StoreDist = false + i++ + } else { + return nil, diceerrors.ErrSyntax + } + case "STOREDIST": + if opts.WithCoord || opts.WithDist || opts.WithHash { + return nil, diceerrors.ErrGeneral("STORE option in GEORADIUS is not compatible with WITHDIST, WITHHASH and WITHCOORD options") + } + if i+1 < len(args) { + opts.Store = args[i+1] + opts.StoreDist = true + i++ + } else { + return nil, diceerrors.ErrSyntax + } + default: + return nil, diceerrors.ErrSyntax + } + } + + return opts, nil +} + +func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { + if len(args) < 4 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEORADIUSBYMEMBER"), + } + } + + key := args[0] + member := args[1] + dist := args[2] + unit := args[3] + + distVal, parseErr := strconv.ParseFloat(dist, 64) + if parseErr != nil { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("need numeric radius"), + } + } + + opts, parseErr := parseGeoRadiusOpts(args[4:]) + if parseErr != nil { + return &EvalResponse{ + Result: nil, + Error: parseErr, + } + } + + obj := store.Get(key) + if obj == nil { + return &EvalResponse{ + Result: clientio.EmptyArray, + Error: nil, + } + } + + ss, err := sortedset.FromObject(obj) + if err != nil { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + } + } + + centerHash, ok := ss.Get(member) + if !ok { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("could not decode requested zset member"), + } + } + + radius, ok := geo.ToMeters(distVal, unit) + if !ok { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrUnsupportedUnit, + } + } + + area, steps := geo.Area(centerHash, radius) + + /* When a huge Radius (in the 5000 km range or more) is used, + * adjacent neighbors can be the same, leading to duplicated + * elements. Skip every range which is the same as the one + * processed previously. */ + var members []string + var hashes []float64 + + anyMax, count := 0, 0 + if opts.CountAny { + anyMax = opts.Count + } + + var lastProcessed uint64 + for _, hash := range area { + if hash == 0 { + continue + } + + if lastProcessed == hash { + continue + } + + hashMin, hashMax := geo.HashMinMax(hash, steps) + rangeMembers, rangeHashes := ss.GetMemberScoresInRange(float64(hashMin), float64(hashMax), count, anyMax) + members = append(members, rangeMembers...) + hashes = append(hashes, rangeHashes...) + count += len(rangeMembers) + lastProcessed = hash + } + + dists := make([]float64, 0, len(members)) + coords := make([][]float64, 0, len(members)) + + centerLat, centerLon := geo.DecodeHash(centerHash) + + for i := range hashes { + msLat, msLon := geo.DecodeHash(hashes[i]) + + dist := geo.GetDistance(centerLon, centerLat, msLon, msLat) + + // Geohash scores are not linear. Therefore, we can sometimes receive results + // which are out of the geographical range and we need to post filter the results here. + if dist > radius { + members[i] = "" + } + + distance, err := geo.ConvertDistance(dist, unit) + if err != nil { + return &EvalResponse{ + Result: nil, + Error: err, + } + } + + dists = append(dists, distance) + coords = append(coords, []float64{msLat, msLon}) + } + + // Sorting is done by distance. Since our output can be dynamic and we can avoid allocating memory + // for each optional output property (hash, dist, coord), we follow an indirect sort approach: + // 1. Save the member inidices. + // 2. Sort the indices based on the distances in ascending or descending order. + // 3. Build the response based on the requested options. + indices := make([]int, len(members)) + for i := range indices { + indices[i] = i + } + + if opts.IsSorted { + if opts.Ascending { + sort.Slice(indices, func(i, j int) bool { + return dists[indices[i]] < dists[indices[j]] + }) + } else { + sort.Slice(indices, func(i, j int) bool { + return dists[indices[i]] > dists[indices[j]] + }) + } + } + + var countVal int + if opts.Count == 0 { + countVal = len(members) + } else { + countVal = opts.Count + } + + if !opts.WithCoord && !opts.WithDist && !opts.WithHash { + response := make([]string, 0, min(len(members), countVal)) + for i := 0; i < cap(response); i++ { + if members[indices[i]] == "" { + continue + } + + response = append(response, members[indices[i]]) + } + + return &EvalResponse{ + Result: response, + Error: nil, + } + } + + response := make([][]interface{}, 0, min(len(members), countVal)) + for i := 0; i < cap(response); i++ { + if members[indices[i]] == "" { + continue + } + + member := []interface{}{} + member = append(member, members[indices[i]]) + if opts.WithDist { + member = append(member, dists[indices[i]]) + } + if opts.WithHash { + member = append(member, hashes[indices[i]]) + } + if opts.WithCoord { + member = append(member, coords[indices[i]]) + } + response = append(response, member) + } + + return &EvalResponse{ + Result: response, + Error: nil, + } +} diff --git a/internal/server/cmd_meta.go b/internal/server/cmd_meta.go index fc9927d4a..8b2e2ca75 100644 --- a/internal/server/cmd_meta.go +++ b/internal/server/cmd_meta.go @@ -458,8 +458,12 @@ var ( Cmd: "GEODIST", Type: SingleShard, } + geoRadiusByMemberCmdMeta = CmdMeta{ + Cmd: "GEORADIUSBYMEMBER", + Type: SingleShard, + } clientCmdMeta = CmdMeta{ - Cmd: "CLIENT", + Cmd: "CLIENT", Type: SingleShard, } latencyCmdMeta = CmdMeta{ @@ -506,7 +510,6 @@ var ( Cmd: "COMMAND|GETKEYSANDFLAGS", Type: SingleShard, } - // Metadata for multishard commands would go here. // These commands require both breakup and gather logic. @@ -637,6 +640,7 @@ func init() { CmdMetaMap["RESTORE"] = restoreCmdMeta CmdMetaMap["GEOADD"] = geoaddCmdMeta CmdMetaMap["GEODIST"] = geodistCmdMeta + CmdMetaMap["GEORADIUSBYMEMBER"] = geoRadiusByMemberCmdMeta CmdMetaMap["CLIENT"] = clientCmdMeta CmdMetaMap["LATENCY"] = latencyCmdMeta CmdMetaMap["FLUSHDB"] = flushDBCmdMeta @@ -649,4 +653,5 @@ func init() { CmdMetaMap["COMMAND|DOCS"] = CmdCommandDocs CmdMetaMap["COMMAND|GETKEYS"] = CmdCommandGetKeys CmdMetaMap["COMMAND|GETKEYSANDFLAGS"] = CmdCommandGetKeysFlags + // Additional commands (multishard, custom) can be added here as needed. }