From 91ca84aef41dd2522765cb7f24a6937dd4ce56f9 Mon Sep 17 00:00:00 2001 From: Ben Date: Mon, 18 Nov 2024 08:04:37 +0100 Subject: [PATCH] internal/eval: implement basic functionality of GEORADIUSBYMEMBER command --- internal/errors/migrated_errors.go | 2 + internal/eval/commands.go | 9 ++ internal/eval/geo/geo.go | 162 +++++++++++++++++++++++++- internal/eval/sortedset/sorted_set.go | 34 +++++- internal/eval/store_eval.go | 90 ++++++++++++++ 5 files changed, 291 insertions(+), 6 deletions(-) diff --git a/internal/errors/migrated_errors.go b/internal/errors/migrated_errors.go index eabb058db3..29525ec3f3 100644 --- a/internal/errors/migrated_errors.go +++ b/internal/errors/migrated_errors.go @@ -33,6 +33,8 @@ var ( ErrInvalidIPAddress = errors.New("invalid IP address") ErrInvalidFingerprint = errors.New("invalid fingerprint") ErrKeyDoesNotExist = errors.New("ERR could not perform this operation on a key that doesn't exist") + 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 787bca20c6..724a4fcb8d 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/geo/geo.go b/internal/eval/geo/geo.go index 6db40eaf8a..1affbd9508 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -13,6 +13,24 @@ const earthRadius float64 = 6372797.560856 // Bit precision for geohash - picked up to match redis const bitPrecision = 52 +const mercatorMax = 20037726.37 + +const ( + minLat = -85.05112878 + maxLat = 85.05112878 + minLon = -180 + maxLon = 180 +) + +type Unit string + +const ( + Meters Unit = "m" + Kilometers Unit = "km" + Miles Unit = "mi" + Feet Unit = "ft" +) + func DegToRad(deg float64) float64 { return math.Pi * deg / 180.0 } @@ -71,16 +89,150 @@ func ConvertDistance( distance float64, unit string, ) (converted float64, err []byte) { - switch unit { - case "m": + 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") } } + +// 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 + } +} + +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) +} + +// 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) ([9]uint64, uint8) { + var result [9]uint64 + + centerLat, centerLon := DecodeHash(centerHash) + + steps := geohashEstimateStepsByRadius(radius, centerLat) + + centerRadiusHash := geohash.EncodeIntWithPrecision(centerLat, centerLon, uint(steps)*2) + + neighbors := geohash.NeighborsIntWithPrecision(centerRadiusHash, uint(steps)*2) + area := geohash.BoundingBoxInt(centerRadiusHash) + + /* 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 := geohash.BoundingBoxInt(neighbors[0]) + east := geohash.BoundingBoxInt(neighbors[2]) + south := geohash.BoundingBoxInt(neighbors[4]) + west := geohash.BoundingBoxInt(neighbors[6]) + + decreaseStep := false + if north.MaxLat < maxLat || south.MinLat < minLat || east.MaxLng < maxLon || west.MinLng < minLon { + decreaseStep = true + } + + if steps > 1 && decreaseStep { + steps-- + centerRadiusHash = geohash.EncodeIntWithPrecision(centerLat, centerLon, uint(steps)*2) + neighbors = geohash.NeighborsIntWithPrecision(centerRadiusHash, uint(steps)*2) + area = geohash.BoundingBoxInt(centerRadiusHash) + } + + // exclude useless areas + if steps >= 2 { + if area.MinLat < minLat { + neighbors[3] = 0 // south east + neighbors[4] = 0 // south + neighbors[5] = 0 // south west + } + + if area.MaxLat > maxLat { + neighbors[0] = 0 // north + neighbors[1] = 0 // north east + neighbors[7] = 0 // north west + } + + if area.MinLng < minLon { + neighbors[5] = 0 // south west + neighbors[6] = 0 // west + neighbors[7] = 0 // north west + } + + if area.MaxLng > maxLon { + neighbors[1] = 0 // north east + neighbors[2] = 0 // east + neighbors[3] = 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) (uint64, uint64) { + min := geohashAlign52Bits(hash, steps) + hash++ + max := geohashAlign52Bits(hash, steps) + return min, max +} + +func geohashAlign52Bits(hash uint64, steps uint8) uint64 { + hash <<= (52 - steps*2) + return hash +} diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index f24a62a931..9afcdbfe3b 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -104,7 +104,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( @@ -251,3 +251,35 @@ func (ss *Set) CountInRange(minVal, maxVal float64) int { return count } + +// GetScoreRange returns a slice of members with scores 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) GetScoreRange( + minScore, maxScore float64, + withScores bool, + reverse bool, +) []string { + var result []string + iterFunc := func(item btree.Item) bool { + ssi := item.(*Item) + if ssi.Score < minScore { + return true + } + if ssi.Score > maxScore { + return false + } + result = append(result, ssi.Member) + if withScores { + scoreStr := strconv.FormatFloat(ssi.Score, 'g', -1, 64) + result = append(result, scoreStr) + } + return true + } + if reverse { + ss.tree.Descend(iterFunc) + } else { + ss.tree.Ascend(iterFunc) + } + return result +} \ No newline at end of file diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index d627c8b6f4..e375b69ad0 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6944,3 +6944,93 @@ func evalCommandDocs(args []string) *EvalResponse { return makeEvalResult(result) } + +func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { + if len(args) < 4 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEODIST"), + } + } + + 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.ErrInvalidFloat, + } + } + + // TODO parse options + // parseGeoRadiusOptions(args[4:]) + + obj := store.Get(key) + if obj == nil { + return &EvalResponse{ + Result: clientio.NIL, + Error: nil, + } + } + + ss, err := sortedset.FromObject(obj) + if err != nil { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + } + } + + memberHash, ok := ss.Get(member) + if !ok { + return &EvalResponse{ + Result: nil, + Error: nil, + } + } + + radius, ok := geo.ToMeters(distVal, unit) + if !ok { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrUnsupportedUnit, + } + } + + area, steps := geo.Area(memberHash, 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 lastProcessed uint64 + for _, hash := range area { + if hash == 0 { + continue + } + + if lastProcessed == hash { + continue + } + + // TODO handle COUNT arg to limit number of returned members + + hashMin, hashMax := geo.HashMinMax(hash, steps) + rangeMembers := ss.GetScoreRange(float64(hashMin), float64(hashMax), false, false) + for _, member := range rangeMembers { + members = append(members, fmt.Sprintf("%q", member)) + } + } + + // TODO handle options + + return &EvalResponse{ + Result: clientio.Encode(members, false), + } +}