From 5c8a10310daeaf6befeb4a2d5e8778b1855ccd6a Mon Sep 17 00:00:00 2001 From: Ben Date: Mon, 18 Nov 2024 08:04:37 +0100 Subject: [PATCH 01/11] 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 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/geo/geo.go b/internal/eval/geo/geo.go index 6db40eaf8..1affbd950 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 5b4435c3f..e1b8365b3 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,35 @@ func DeserializeSortedSet(buf *bytes.Reader) (*Set, error) { return ss, nil } + +// 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 +} diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index c80ad9d63..198275c02 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6772,3 +6772,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), + } +} From 5ba19c37a0672371b75e87fe547c656524108d08 Mon Sep 17 00:00:00 2001 From: Ben Date: Fri, 22 Nov 2024 08:37:43 +0100 Subject: [PATCH 02/11] internal/eval: add support for WITHDIST, WITHHASH, WITHCOORD and COUNT [ANY] options in GEORADIUSBYMEMBER cmd --- internal/clientio/resp.go | 19 +++ internal/eval/sortedset/sorted_set.go | 31 ++-- internal/eval/store_eval.go | 225 ++++++++++++++++++++++++-- internal/server/cmd_meta.go | 9 +- 4 files changed, 248 insertions(+), 36 deletions(-) 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/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index e1b8365b3..d1e8f0516 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -311,14 +311,11 @@ func DeserializeSortedSet(buf *bytes.Reader) (*Set, error) { } // 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 +func (ss *Set) GetMemberScoresInRange(minScore, maxScore float64, count, max int) ([]string, []float64) { + var members []string + var scores []float64 + iterFunc := func(item btree.Item) bool { ssi := item.(*Item) if ssi.Score < minScore { @@ -327,17 +324,17 @@ func (ss *Set) GetScoreRange( 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) + members = append(members, ssi.Member) + scores = append(scores, ssi.Score) + count++ + + if max > 0 && count == max { + return false } + return true } - if reverse { - ss.tree.Descend(iterFunc) - } else { - ss.tree.Ascend(iterFunc) - } - return result + + ss.tree.Ascend(iterFunc) + return members, scores } diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 198275c02..e0f107662 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6773,6 +6773,90 @@ 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 + StoreDist string +} + +func parseGeoRadiusOpts(args []string) (*geoRadiusOpts, error) { + opts := &geoRadiusOpts{ + Ascending: true, // Default to ascending order if sorted + } + + for i := 0; i < len(args); i++ { + param := strings.ToUpper(args[i]) + + switch param { + case "WITHDIST": + opts.WithDist = true + case "WITHCOORD": + opts.WithCoord = true + case "WITHHASH": + opts.WithHash = true + case "COUNT": + + // TODO validate this logic + + if i+1 >= len(args) { + return nil, fmt.Errorf("ERR syntax error") + } + + count, err := strconv.Atoi(args[i+1]) + if err != nil { + return nil, fmt.Errorf("ERR value is not an integer or out of range") + } + if count <= 0 { + return nil, fmt.Errorf("ERR COUNT must be > 0") + } + opts.Count = count + i++ + + // Check for ANY option after COUNT + if i+1 < len(args) && strings.ToUpper(args[i+1]) == "ANY" { + opts.CountAny = true + i++ + } + case "ASC": + opts.IsSorted = true + opts.Ascending = true + + case "DESC": + opts.IsSorted = true + opts.Ascending = false + + case "STORE": + if i+1 >= len(args) { + return nil, fmt.Errorf("STORE option requires a key name") + } + opts.Store = args[i+1] + i++ + + case "STOREDIST": + if i+1 >= len(args) { + return nil, fmt.Errorf("STOREDIST option requires a key name") + } + opts.StoreDist = args[i+1] + i++ + + default: + return nil, fmt.Errorf("unknown parameter: %s", args[i]) + } + } + + if opts.Store != "" && opts.StoreDist != "" { + return nil, fmt.Errorf("STORE and STOREDIST are mutually exclusive") + } + + return opts, nil +} + func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { if len(args) < 4 { return &EvalResponse{ @@ -6789,31 +6873,33 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { distVal, parseErr := strconv.ParseFloat(dist, 64) if parseErr != nil { return &EvalResponse{ - Result: nil, - Error: diceerrors.ErrInvalidFloat, + Error: diceerrors.ErrInvalidFloat, } } - // TODO parse options - // parseGeoRadiusOptions(args[4:]) + 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.NIL, - Error: nil, } } ss, err := sortedset.FromObject(obj) if err != nil { return &EvalResponse{ - Result: nil, - Error: diceerrors.ErrWrongTypeOperation, + Error: diceerrors.ErrWrongTypeOperation, } } - memberHash, ok := ss.Get(member) + centerHash, ok := ss.Get(member) if !ok { return &EvalResponse{ Result: nil, @@ -6829,14 +6915,20 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { } } - area, steps := geo.Area(memberHash, radius) + 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 { @@ -6847,18 +6939,117 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { 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)) + rangeMembers, rangeHashes := ss.GetMemberScoresInRange(float64(hashMin), float64(hashMax), count, anyMax) + members = append(members, rangeMembers...) + hashes = append(hashes, rangeHashes...) + } + + dists := make([]float64, 0, len(members)) + coords := make([][]float64, 0, len(members)) + + centerLat, centerLon := geo.DecodeHash(centerHash) + + if opts.IsSorted || opts.WithDist || opts.WithCoord { + for i := range hashes { + msLat, msLon := geo.DecodeHash(hashes[i]) + + if opts.WithDist || opts.IsSorted { + dist := geo.GetDistance(centerLon, centerLat, msLon, msLat) + dists = append(dists, dist) + } + + if opts.WithCoord { + 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]] + }) } } - // TODO handle options + optCount := 0 + if opts.WithDist { + optCount++ + } + + if opts.WithHash { + optCount++ + } + + if opts.WithCoord { + optCount++ + } + + max := opts.Count + if max > len(members) { + max = len(members) + } + + if optCount == 0 { + response := make([]string, len(members)) + for i := range members { + response[i] = members[indices[i]] + } + + if max > 0 { + response = response[:max] + } + + return &EvalResponse{ + Result: clientio.Encode(response, false), + } + } + + response := make([][]interface{}, len(members)) + for i := range members { + item := make([]any, optCount+1) + item[0] = members[i] + + itemIdx := 1 + + if opts.WithDist { + item[itemIdx] = dists[i] + itemIdx++ + } + + if opts.WithHash { + item[itemIdx] = hashes[i] + itemIdx++ + } + + if opts.WithCoord { + item[itemIdx] = coords[i] + itemIdx++ + } + + response[indices[i]] = item + } + + if max > 0 { + response = response[:max] + } return &EvalResponse{ - Result: clientio.Encode(members, false), + Result: clientio.Encode(response, false), } } 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. } From d577da58f877434ebd4fa903f54569b2a2355e22 Mon Sep 17 00:00:00 2001 From: c-harish Date: Sat, 23 Nov 2024 10:54:20 +0530 Subject: [PATCH 03/11] feat: add georadius optional arguments validation --- internal/eval/store_eval.go | 102 +++++++++++++++++------------------- 1 file changed, 48 insertions(+), 54 deletions(-) diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index e0f107662..3c6f5ee3b 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6777,83 +6777,77 @@ 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 - StoreDist string + 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{ - Ascending: true, // Default to ascending order if sorted - } + opts := &geoRadiusOpts{} for i := 0; i < len(args); i++ { - param := strings.ToUpper(args[i]) + option := strings.ToUpper(args[i]) - switch param { - case "WITHDIST": - opts.WithDist = true + switch option { case "WITHCOORD": opts.WithCoord = true + case "WITHDIST": + opts.WithDist = true case "WITHHASH": opts.WithHash = true - case "COUNT": - - // TODO validate this logic - - if i+1 >= len(args) { - return nil, fmt.Errorf("ERR syntax error") - } - - count, err := strconv.Atoi(args[i+1]) - if err != nil { - return nil, fmt.Errorf("ERR value is not an integer or out of range") - } - if count <= 0 { - return nil, fmt.Errorf("ERR COUNT must be > 0") - } - opts.Count = count - i++ - - // Check for ANY option after COUNT - if i+1 < len(args) && strings.ToUpper(args[i+1]) == "ANY" { - opts.CountAny = true - i++ - } 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.ErrGeneral("syntax error") + } + case "ANY": + return nil, diceerrors.ErrGeneral("the ANY argument requires COUNT argument") case "STORE": - if i+1 >= len(args) { - return nil, fmt.Errorf("STORE option requires a key name") + 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.ErrGeneral("syntax error") } - opts.Store = args[i+1] - i++ - case "STOREDIST": - if i+1 >= len(args) { - return nil, fmt.Errorf("STOREDIST option requires a key name") + 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.ErrGeneral("syntax error") } - opts.StoreDist = args[i+1] - i++ - default: - return nil, fmt.Errorf("unknown parameter: %s", args[i]) + return nil, diceerrors.ErrGeneral("syntax error") } } - if opts.Store != "" && opts.StoreDist != "" { - return nil, fmt.Errorf("STORE and STOREDIST are mutually exclusive") - } - return opts, nil } @@ -6881,7 +6875,7 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { if parseErr != nil { return &EvalResponse{ Result: nil, - Error: parseErr, + Error: parseErr, } } From 81b49915ffa0df38fef08c6d7f369542d81d617a Mon Sep 17 00:00:00 2001 From: c-harish Date: Sat, 23 Nov 2024 11:43:17 +0530 Subject: [PATCH 04/11] fix: update return responses --- internal/eval/store_eval.go | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 3c6f5ee3b..90f83d312 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6817,7 +6817,7 @@ func parseGeoRadiusOpts(args []string) (*geoRadiusOpts, error) { } i++ } else { - return nil, diceerrors.ErrGeneral("syntax error") + return nil, diceerrors.ErrSyntax } case "ANY": return nil, diceerrors.ErrGeneral("the ANY argument requires COUNT argument") @@ -6830,7 +6830,7 @@ func parseGeoRadiusOpts(args []string) (*geoRadiusOpts, error) { opts.StoreDist = false i++ } else { - return nil, diceerrors.ErrGeneral("syntax error") + return nil, diceerrors.ErrSyntax } case "STOREDIST": if opts.WithCoord || opts.WithDist || opts.WithHash { @@ -6841,10 +6841,10 @@ func parseGeoRadiusOpts(args []string) (*geoRadiusOpts, error) { opts.StoreDist = true i++ } else { - return nil, diceerrors.ErrGeneral("syntax error") + return nil, diceerrors.ErrSyntax } default: - return nil, diceerrors.ErrGeneral("syntax error") + return nil, diceerrors.ErrSyntax } } @@ -6855,7 +6855,7 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { if len(args) < 4 { return &EvalResponse{ Result: nil, - Error: diceerrors.ErrWrongArgumentCount("GEODIST"), + Error: diceerrors.ErrWrongArgumentCount("GEORADIUSBYMEMBER"), } } @@ -6867,7 +6867,8 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { distVal, parseErr := strconv.ParseFloat(dist, 64) if parseErr != nil { return &EvalResponse{ - Error: diceerrors.ErrInvalidFloat, + Result: nil, + Error: diceerrors.ErrGeneral("need numeric radius"), } } @@ -6882,14 +6883,16 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { obj := store.Get(key) if obj == nil { return &EvalResponse{ - Result: clientio.NIL, + Result: clientio.EmptyArray, + Error: nil, } } ss, err := sortedset.FromObject(obj) if err != nil { return &EvalResponse{ - Error: diceerrors.ErrWrongTypeOperation, + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, } } @@ -6897,7 +6900,7 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { if !ok { return &EvalResponse{ Result: nil, - Error: nil, + Error: diceerrors.ErrGeneral("could not decode requested zset member"), } } From b4b9a5906bc84378bf81e47a804e72c36633d1e6 Mon Sep 17 00:00:00 2001 From: c-harish Date: Sat, 23 Nov 2024 14:22:29 +0530 Subject: [PATCH 05/11] fix: fix response structure --- internal/eval/geo/geo.go | 9 ++-- internal/eval/store_eval.go | 85 ++++++++++++------------------------- 2 files changed, 31 insertions(+), 63 deletions(-) diff --git a/internal/eval/geo/geo.go b/internal/eval/geo/geo.go index 1affbd950..872b15b6e 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -3,7 +3,7 @@ package geo import ( "math" - "github.com/dicedb/dice/internal/errors" + diceerrors "github.com/dicedb/dice/internal/errors" "github.com/mmcloughlin/geohash" ) @@ -85,10 +85,7 @@ func DecodeHash(hash float64) (lat, lon float64) { } // ConvertDistance converts a distance from meters to the desired unit -func ConvertDistance( - distance float64, - unit string, -) (converted float64, err []byte) { +func ConvertDistance(distance float64, unit string) (float64, error) { switch Unit(unit) { case Meters: return distance, nil @@ -99,7 +96,7 @@ func ConvertDistance( 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 } } diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 90f83d312..488ded24c 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, } } @@ -6940,6 +6940,8 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { 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)) @@ -6953,7 +6955,14 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { if opts.WithDist || opts.IsSorted { dist := geo.GetDistance(centerLon, centerLat, msLon, msLat) - dists = append(dists, dist) + distance, err := geo.ConvertDistance(dist, unit) + if err != nil { + return &EvalResponse{ + Result: nil, + Error: err, + } + } + dists = append(dists, distance) } if opts.WithCoord { @@ -6984,69 +6993,31 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { } } - optCount := 0 - if opts.WithDist { - optCount++ - } - - if opts.WithHash { - optCount++ - } - - if opts.WithCoord { - optCount++ - } - - max := opts.Count - if max > len(members) { - max = len(members) - } - - if optCount == 0 { - response := make([]string, len(members)) - for i := range members { - response[i] = members[indices[i]] - } - - if max > 0 { - response = response[:max] - } - - return &EvalResponse{ - Result: clientio.Encode(response, false), - } + var countVal int + if opts.Count == 0 { + countVal = len(members) + } else { + countVal = opts.Count } - response := make([][]interface{}, len(members)) - for i := range members { - item := make([]any, optCount+1) - item[0] = members[i] - - itemIdx := 1 - + response := make([][]interface{}, 0, min(len(members), countVal)) + for i := 0; i < cap(response); i++ { + member := []interface{}{} + member = append(member, members[indices[i]]) if opts.WithDist { - item[itemIdx] = dists[i] - itemIdx++ + member = append(member, dists[indices[i]]) } - if opts.WithHash { - item[itemIdx] = hashes[i] - itemIdx++ + member = append(member, hashes[indices[i]]) } - if opts.WithCoord { - item[itemIdx] = coords[i] - itemIdx++ + member = append(member, coords[indices[i]]) } - - response[indices[i]] = item - } - - if max > 0 { - response = response[:max] + response = append(response, member) } return &EvalResponse{ - Result: clientio.Encode(response, false), + Result: response, + Error: nil, } } From da37645b298203c336b7cd00240312309384ad8d Mon Sep 17 00:00:00 2001 From: Ben Date: Wed, 27 Nov 2024 07:49:22 +0100 Subject: [PATCH 06/11] internal/eval: replace geohash library with redis-like implementation and finalize GEORADIUSBYMEMBER command --- go.mod | 1 - go.sum | 2 - internal/eval/geo/geo.go | 303 ++++++++++++++++++++++---- internal/eval/sortedset/sorted_set.go | 4 +- internal/eval/store_eval.go | 37 ++-- 5 files changed, 281 insertions(+), 66 deletions(-) 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/eval/geo/geo.go b/internal/eval/geo/geo.go index 872b15b6e..a6e89cdd2 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -4,22 +4,23 @@ import ( "math" diceerrors "github.com/dicedb/dice/internal/errors" - "github.com/mmcloughlin/geohash" ) // 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 ( - minLat = -85.05112878 - maxLat = 85.05112878 - minLon = -180 - maxLon = 180 + /* 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 @@ -31,20 +32,18 @@ const ( 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) @@ -62,29 +61,73 @@ 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 float64, 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 +// ConvertDistance converts a distance from meters to the desired unit. func ConvertDistance(distance float64, unit string) (float64, error) { switch Unit(unit) { case Meters: @@ -100,7 +143,7 @@ func ConvertDistance(distance float64, unit string) (float64, error) { } } -// ToMeters converts a distance and its unit to meters +// ToMeters converts a distance and its unit to meters. func ToMeters(distance float64, unit string) (float64, bool) { switch Unit(unit) { case Meters: @@ -116,6 +159,7 @@ func ToMeters(distance float64, unit string) (float64, bool) { } } +// 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 @@ -149,63 +193,89 @@ func geohashEstimateStepsByRadius(radius, lat float64) uint8 { return uint8(step) } +// boundingBox returns the bounding box for a given latitude, longitude and radius. +func boundingBox(lat, lon, radius float64) (float64, float64, float64, 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) ([9]uint64, uint8) { var result [9]uint64 - centerLat, centerLon := DecodeHash(centerHash) - + centerLat, centerLon := decodeHash(uint64(centerHash), maxSteps) + minLon, minLat, maxLon, maxLat := boundingBox(centerLat, centerLon, radius) steps := geohashEstimateStepsByRadius(radius, centerLat) + centerRadiusHash := encodeHash(centerLon, centerLat, steps) - centerRadiusHash := geohash.EncodeIntWithPrecision(centerLat, centerLon, uint(steps)*2) - - neighbors := geohash.NeighborsIntWithPrecision(centerRadiusHash, uint(steps)*2) - area := geohash.BoundingBoxInt(centerRadiusHash) + neighbors := geohashNeighbors(uint64(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 := geohash.BoundingBoxInt(neighbors[0]) - east := geohash.BoundingBoxInt(neighbors[2]) - south := geohash.BoundingBoxInt(neighbors[4]) - west := geohash.BoundingBoxInt(neighbors[6]) + north := areaBySteps(neighbors[0], steps) + south := areaBySteps(neighbors[4], steps) + east := areaBySteps(neighbors[2], steps) + west := areaBySteps(neighbors[6], steps) decreaseStep := false - if north.MaxLat < maxLat || south.MinLat < minLat || east.MaxLng < maxLon || west.MinLng < minLon { + 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 = geohash.EncodeIntWithPrecision(centerLat, centerLon, uint(steps)*2) - neighbors = geohash.NeighborsIntWithPrecision(centerRadiusHash, uint(steps)*2) - area = geohash.BoundingBoxInt(centerRadiusHash) + centerRadiusHash = encodeHash(centerLat, centerLon, steps) + neighbors = geohashNeighbors(centerRadiusHash, steps) + area = areaBySteps(centerRadiusHash, steps) } // exclude useless areas if steps >= 2 { - if area.MinLat < minLat { + if area.Lat.Min < minLat { neighbors[3] = 0 // south east neighbors[4] = 0 // south neighbors[5] = 0 // south west } - if area.MaxLat > maxLat { + if area.Lat.Max > maxLat { neighbors[0] = 0 // north neighbors[1] = 0 // north east neighbors[7] = 0 // north west } - if area.MinLng < minLon { + if area.Lon.Min < minLon { neighbors[5] = 0 // south west neighbors[6] = 0 // west neighbors[7] = 0 // north west } - if area.MaxLng > maxLon { + if area.Lon.Max > maxLon { neighbors[1] = 0 // north east neighbors[2] = 0 // east neighbors[3] = 0 // south east @@ -223,13 +293,156 @@ func Area(centerHash, radius float64) ([9]uint64, uint8) { // 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) + min := align52Bits(hash, steps) hash++ - max := geohashAlign52Bits(hash, steps) + max := align52Bits(hash, steps) return min, max } -func geohashAlign52Bits(hash uint64, steps uint8) uint64 { +// 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 = x + uint64(zz+1) + } else { + x = x | uint64(zz) + x = 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 = y + (zz + 1) + } else { + y = y | zz + y = 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] = geohashMoveX(geohashMoveY(hash, steps, 1), steps, 1) // North-East + neighbors[2] = geohashMoveX(hash, steps, 1) // East + neighbors[3] = geohashMoveX(geohashMoveY(hash, steps, -1), steps, 1) // South-East + neighbors[4] = geohashMoveY(hash, steps, -1) // South + neighbors[5] = geohashMoveX(geohashMoveY(hash, steps, -1), steps, -1) // South-West + neighbors[6] = geohashMoveX(hash, steps, -1) // West + neighbors[7] = geohashMoveX(geohashMoveY(hash, steps, 1), steps, -1) // North-West + + return neighbors +} diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index d1e8f0516..c5d927823 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -321,14 +321,14 @@ func (ss *Set) GetMemberScoresInRange(minScore, maxScore float64, count, max int if ssi.Score < minScore { return true } - if ssi.Score > maxScore { + if ssi.Score >= maxScore { return false } members = append(members, ssi.Member) scores = append(scores, ssi.Score) count++ - if max > 0 && count == max { + if max > 0 && count >= max { return false } diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 488ded24c..4953e34d4 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -6949,26 +6949,27 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { centerLat, centerLon := geo.DecodeHash(centerHash) - if opts.IsSorted || opts.WithDist || opts.WithCoord { - for i := range hashes { - msLat, msLon := geo.DecodeHash(hashes[i]) + for i := range hashes { + msLat, msLon := geo.DecodeHash(hashes[i]) - if opts.WithDist || opts.IsSorted { - dist := geo.GetDistance(centerLon, centerLat, msLon, msLat) - distance, err := geo.ConvertDistance(dist, unit) - if err != nil { - return &EvalResponse{ - Result: nil, - Error: err, - } - } - dists = append(dists, distance) - } + 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] = "" + } - if opts.WithCoord { - coords = append(coords, []float64{msLat, msLon}) + 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 @@ -7002,6 +7003,10 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { 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 { From 8e8a4050a1e7c4f1c86a0ba47f64c63adfa6db08 Mon Sep 17 00:00:00 2001 From: Ben Date: Thu, 28 Nov 2024 07:44:47 +0100 Subject: [PATCH 07/11] internal/eval: fix response structure for GEORADIUSBYMEMBER when no parameters are set --- internal/eval/store_eval.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index 4953e34d4..25dd6effd 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -7001,6 +7001,22 @@ func evalGEORADIUSBYMEMBER(args []string, store *dstore.Store) *EvalResponse { 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]] == "" { From 375143beb297224366eb881fc894bb9ac91b613d Mon Sep 17 00:00:00 2001 From: Ben Date: Fri, 29 Nov 2024 08:26:27 +0100 Subject: [PATCH 08/11] internal/eval: add unit tests for GEORADIUSBYMEMBER and fix small bugs in processing neighbors and calculating hashes on step reduction --- internal/eval/eval_test.go | 142 ++++++++++++++++++++++++++++++++++++- internal/eval/geo/geo.go | 38 +++++----- 2 files changed, 159 insertions(+), 21 deletions(-) diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index f1a629a6e..2c3eaa2d0 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -132,6 +132,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 +8296,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 +8307,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 +9223,140 @@ 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, object.ObjEncodingRaw)) + }, + 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) +} \ No newline at end of file diff --git a/internal/eval/geo/geo.go b/internal/eval/geo/geo.go index a6e89cdd2..1baa33990 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -239,9 +239,9 @@ func Area(centerHash, radius float64) ([9]uint64, uint8) { * north / south / west / east square is too near to the search area * to cover everything. */ north := areaBySteps(neighbors[0], steps) - south := areaBySteps(neighbors[4], steps) + south := areaBySteps(neighbors[1], steps) east := areaBySteps(neighbors[2], steps) - west := areaBySteps(neighbors[6], 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 { @@ -250,7 +250,7 @@ func Area(centerHash, radius float64) ([9]uint64, uint8) { if steps > 1 && decreaseStep { steps-- - centerRadiusHash = encodeHash(centerLat, centerLon, steps) + centerRadiusHash = encodeHash(centerLon, centerLat, steps) neighbors = geohashNeighbors(centerRadiusHash, steps) area = areaBySteps(centerRadiusHash, steps) } @@ -258,27 +258,27 @@ func Area(centerHash, radius float64) ([9]uint64, uint8) { // exclude useless areas if steps >= 2 { if area.Lat.Min < minLat { - neighbors[3] = 0 // south east - neighbors[4] = 0 // south - neighbors[5] = 0 // south west + neighbors[6] = 0 // south east + neighbors[1] = 0 // south + neighbors[7] = 0 // south west } if area.Lat.Max > maxLat { neighbors[0] = 0 // north - neighbors[1] = 0 // north east - neighbors[7] = 0 // north west + neighbors[4] = 0 // north east + neighbors[5] = 0 // north west } if area.Lon.Min < minLon { - neighbors[5] = 0 // south west - neighbors[6] = 0 // west - neighbors[7] = 0 // north west + neighbors[7] = 0 // south west + neighbors[3] = 0 // west + neighbors[5] = 0 // north west } if area.Lon.Max > maxLon { - neighbors[1] = 0 // north east + neighbors[4] = 0 // north east neighbors[2] = 0 // east - neighbors[3] = 0 // south east + neighbors[6] = 0 // south east } } @@ -436,13 +436,13 @@ func geohashNeighbors(hash uint64, steps uint8) [8]uint64 { neighbors := [8]uint64{} neighbors[0] = geohashMoveY(hash, steps, 1) // North - neighbors[1] = geohashMoveX(geohashMoveY(hash, steps, 1), steps, 1) // North-East + neighbors[1] = geohashMoveY(hash, steps, -1) // South neighbors[2] = geohashMoveX(hash, steps, 1) // East - neighbors[3] = geohashMoveX(geohashMoveY(hash, steps, -1), steps, 1) // South-East - neighbors[4] = geohashMoveY(hash, steps, -1) // South - neighbors[5] = geohashMoveX(geohashMoveY(hash, steps, -1), steps, -1) // South-West - neighbors[6] = geohashMoveX(hash, steps, -1) // West - neighbors[7] = geohashMoveX(geohashMoveY(hash, steps, 1), steps, -1) // North-West + 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 } From 1c4ac7ecd21799d7fb324853eb794b8821ff84f4 Mon Sep 17 00:00:00 2001 From: Ben Date: Sun, 1 Dec 2024 19:32:27 +0100 Subject: [PATCH 09/11] internal/eval: add benchmarks for GEORADIUSBYMEMBER --- internal/eval/eval_test.go | 430 +++++++++++++++++++++++++------------ 1 file changed, 296 insertions(+), 134 deletions(-) diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 2c3eaa2d0..63327fa9b 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" @@ -9225,138 +9226,299 @@ func testEvalLRANGE(t *testing.T, store *dstore.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, object.ObjEncodingRaw)) - }, - 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, - }, - }, + 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, object.ObjEncodingRaw)) + }, + 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) -} \ No newline at end of file + 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) + } + } + }) + } +} From 275bcc6f1b4c2fcdd87d5fca9c85d4e2c1ce3255 Mon Sep 17 00:00:00 2001 From: Ben Date: Mon, 2 Dec 2024 07:40:10 +0100 Subject: [PATCH 10/11] internal/eval: fix linter issues for new GEORADIUSBYMEMBER command --- internal/eval/geo/geo.go | 38 +++++++++++++-------------- internal/eval/sortedset/sorted_set.go | 6 ++--- 2 files changed, 20 insertions(+), 24 deletions(-) diff --git a/internal/eval/geo/geo.go b/internal/eval/geo/geo.go index 1baa33990..de21193d6 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -91,7 +91,7 @@ func DecodeHash(hash float64) (lat, lon float64) { } // decodeHash decodes the geohash into latitude and longitude with the specified number of steps. -func decodeHash(hash uint64, steps uint8) (lat float64, lon float64) { +func decodeHash(hash uint64, steps uint8) (lat, lon float64) { hashSep := deinterleave64(hash) latScale := globalMaxLat - globalMinLat @@ -194,7 +194,7 @@ func geohashEstimateStepsByRadius(radius, lat float64) uint8 { } // boundingBox returns the bounding box for a given latitude, longitude and radius. -func boundingBox(lat, lon, radius float64) (float64, float64, float64, float64) { +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))) @@ -204,30 +204,28 @@ func boundingBox(lat, lon, radius float64) (float64, float64, float64, float64) isSouthernHemisphere = true } - minLon := lon - lonDeltaTop + minLon = lon - lonDeltaTop if isSouthernHemisphere { minLon = lon - lonDeltaBottom } - maxLon := lon + lonDeltaTop + maxLon = lon + lonDeltaTop if isSouthernHemisphere { maxLon = lon + lonDeltaBottom } - minLat := lat - latDelta - maxLat := lat + latDelta + 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) ([9]uint64, uint8) { - var result [9]uint64 - +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) + steps = geohashEstimateStepsByRadius(radius, centerLat) centerRadiusHash := encodeHash(centerLon, centerLat, steps) neighbors := geohashNeighbors(uint64(centerRadiusHash), steps) @@ -292,11 +290,11 @@ func Area(centerHash, radius float64) ([9]uint64, uint8) { // 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 := align52Bits(hash, steps) +func HashMinMax(hash uint64, steps uint8) (minHash uint64, maxHash uint64) { + minHash = align52Bits(hash, steps) hash++ - max := align52Bits(hash, steps) - return min, max + maxHash = align52Bits(hash, steps) + return minHash, maxHash } // align52Bits aligns the hash to 52 bits. @@ -403,10 +401,10 @@ func geohashMoveX(hash uint64, steps uint8, d int8) uint64 { zz := 0x5555555555555555 >> (64 - steps*2) if d > 0 { - x = x + uint64(zz+1) + x += uint64(zz + 1) } else { - x = x | uint64(zz) - x = x - uint64(zz+1) + x |= uint64(zz) + x -= uint64(zz + 1) } x &= (0xaaaaaaaaaaaaaaaa >> (64 - steps*2)) @@ -421,10 +419,10 @@ func geohashMoveY(hash uint64, steps uint8, d int8) uint64 { zz := uint64(0xaaaaaaaaaaaaaaaa) >> (64 - steps*2) if d > 0 { - y = y + (zz + 1) + y += (zz + 1) } else { - y = y | zz - y = y - (zz + 1) + y |= zz + y -= (zz + 1) } y &= (0x5555555555555555 >> (64 - steps*2)) diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index c5d927823..715914543 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -312,9 +312,7 @@ func DeserializeSortedSet(buf *bytes.Reader) (*Set, error) { // 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, max int) ([]string, []float64) { - var members []string - var scores []float64 +func (ss *Set) GetMemberScoresInRange(minScore, maxScore float64, count, maxCount int) (members []string, scores []float64) { iterFunc := func(item btree.Item) bool { ssi := item.(*Item) @@ -328,7 +326,7 @@ func (ss *Set) GetMemberScoresInRange(minScore, maxScore float64, count, max int scores = append(scores, ssi.Score) count++ - if max > 0 && count >= max { + if maxCount > 0 && count >= maxCount { return false } From 203f35dc2e40fa82cd2d838a7f3580cd30395714 Mon Sep 17 00:00:00 2001 From: Ben Date: Sun, 8 Dec 2024 19:21:52 +0100 Subject: [PATCH 11/11] internal/eval: fix tests for GEORADIUSBYMEMBER command and linter issues --- internal/eval/eval_test.go | 2 +- internal/eval/geo/geo.go | 4 ++-- internal/eval/sortedset/sorted_set.go | 1 - 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 63327fa9b..73be73400 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -9250,7 +9250,7 @@ func testEvalGEORADIUSBYMEMBER(t *testing.T, store *dstore.Store) { }, "GEORADIUSBYMEMBER wrong type operation": { setup: func() { - store.Put("wrongtype", store.NewObj("string_value", -1, object.ObjTypeString, object.ObjEncodingRaw)) + store.Put("wrongtype", store.NewObj("string_value", -1, object.ObjTypeString)) }, input: []string{"wrongtype", "wtc one", "7", "km"}, migratedOutput: EvalResponse{ diff --git a/internal/eval/geo/geo.go b/internal/eval/geo/geo.go index de21193d6..0cfd2e776 100644 --- a/internal/eval/geo/geo.go +++ b/internal/eval/geo/geo.go @@ -228,7 +228,7 @@ func Area(centerHash, radius float64) (result [9]uint64, steps uint8) { steps = geohashEstimateStepsByRadius(radius, centerLat) centerRadiusHash := encodeHash(centerLon, centerLat, steps) - neighbors := geohashNeighbors(uint64(centerRadiusHash), steps) + neighbors := geohashNeighbors(centerRadiusHash, steps) area := areaBySteps(centerRadiusHash, steps) /* Check if the step is enough at the limits of the covered area. @@ -290,7 +290,7 @@ func Area(centerHash, radius float64) (result [9]uint64, steps uint8) { // 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 uint64, maxHash uint64) { +func HashMinMax(hash uint64, steps uint8) (minHash, maxHash uint64) { minHash = align52Bits(hash, steps) hash++ maxHash = align52Bits(hash, steps) diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go index 715914543..cde8e575f 100644 --- a/internal/eval/sortedset/sorted_set.go +++ b/internal/eval/sortedset/sorted_set.go @@ -313,7 +313,6 @@ func DeserializeSortedSet(buf *bytes.Reader) (*Set, error) { // 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 {