internal/eval: implement basic functionality of GEORADIUSBYMEMBER com…

…mand

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 {

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

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
+}

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
+}

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),
+	}
+}