From 99cb709773fe15dcf7d374adce412c8d9c0e1f1e Mon Sep 17 00:00:00 2001 From: Uddeshya Singh Date: Mon, 24 Jul 2023 12:42:07 +0530 Subject: [PATCH] Add support for GEOSEARCH and GEOSEARCHSTORE (#1533) --- src/commands/cmd_geo.cc | 258 +++++++++++++++++++++++++++++- src/types/geohash.cc | 80 +++++---- src/types/geohash.h | 20 ++- src/types/redis_geo.cc | 110 ++++++++++--- src/types/redis_geo.h | 20 ++- tests/gocase/unit/geo/geo_test.go | 64 ++++++++ 6 files changed, 478 insertions(+), 74 deletions(-) diff --git a/src/commands/cmd_geo.cc b/src/commands/cmd_geo.cc index 62dab23c338..0b732833246 100644 --- a/src/commands/cmd_geo.cc +++ b/src/commands/cmd_geo.cc @@ -18,9 +18,11 @@ * */ +#include "command_parser.h" #include "commander.h" #include "error_constants.h" #include "server/server.h" +#include "types/geohash.h" #include "types/redis_geo.h" namespace redis { @@ -52,10 +54,14 @@ class CommandGeoBase : public Commander { *longitude = *long_stat; *latitude = *lat_stat; + return ValidateLongLat(longitude, latitude); + } + + static Status ValidateLongLat(double *longitude, double *latitude) { if (*longitude < GEO_LONG_MIN || *longitude > GEO_LONG_MAX || *latitude < GEO_LAT_MIN || *latitude > GEO_LAT_MAX) { - return {Status::RedisParseErr, "invalid longitude,latitude pair " + longitude_para + "," + latitude_para}; + return {Status::RedisParseErr, + "invalid longitude,latitude pair " + std::to_string(*longitude) + "," + std::to_string(*latitude)}; } - return Status::OK(); } @@ -355,6 +361,249 @@ class CommandGeoRadius : public CommandGeoBase { double latitude_ = 0; }; +class CommandGeoSearch : public CommandGeoBase { + public: + CommandGeoSearch() : CommandGeoBase() {} + + Status Parse(const std::vector &args) override { + CommandParser parser(args, 1); + key_ = GET_OR_RET(parser.TakeStr()); + + while (parser.Good()) { + if (parser.EatEqICase("frommember")) { + auto s = setOriginType(kMember); + if (!s.IsOK()) return s; + + member_ = GET_OR_RET(parser.TakeStr()); + } else if (parser.EatEqICase("fromlonlat")) { + auto s = setOriginType(kLongLat); + if (!s.IsOK()) return s; + + longitude_ = GET_OR_RET(parser.TakeFloat()); + latitude_ = GET_OR_RET(parser.TakeFloat()); + s = ValidateLongLat(&longitude_, &latitude_); + if (!s.IsOK()) return s; + } else if (parser.EatEqICase("byradius")) { + auto s = setShapeType(kGeoShapeTypeCircular); + if (!s.IsOK()) return s; + radius_ = GET_OR_RET(parser.TakeFloat()); + std::string distance_raw = GET_OR_RET(parser.TakeStr()); + s = ParseDistanceUnit(distance_raw); + if (!s.IsOK()) return s; + } else if (parser.EatEqICase("bybox")) { + auto s = setShapeType(kGeoShapeTypeRectangular); + if (!s.IsOK()) return s; + width_ = GET_OR_RET(parser.TakeFloat()); + height_ = GET_OR_RET(parser.TakeFloat()); + std::string distance_raw = GET_OR_RET(parser.TakeStr()); + s = ParseDistanceUnit(distance_raw); + if (!s.IsOK()) return s; + } else if (parser.EatEqICase("asc") && sort_ == kSortNone) { + sort_ = kSortASC; + } else if (parser.EatEqICase("desc") && sort_ == kSortNone) { + sort_ = kSortDESC; + } else if (parser.EatEqICase("count")) { + count_ = GET_OR_RET(parser.TakeInt(NumericRange{1, std::numeric_limits::max()})); + } else if (parser.EatEqICase("withcoord")) { + with_coord_ = true; + } else if (parser.EatEqICase("withdist")) { + with_dist_ = true; + } else if (parser.EatEqICase("withhash")) { + with_hash_ = true; + } else { + return {Status::RedisParseErr, "Invalid argument given"}; + } + } + + if (member_ != "" && longitude_ != 0 && latitude_ != 0) { + return {Status::RedisParseErr, "please use only one of FROMMEMBER or FROMLONLAT"}; + } + + auto s = createGeoShape(); + if (!s.IsOK()) { + return s; + } + return Commander::Parse(args); + } + + Status Execute(Server *svr, Connection *conn, std::string *output) override { + std::vector geo_points; + redis::Geo geo_db(svr->storage, conn->GetNamespace()); + + auto s = geo_db.Search(args_[1], geo_shape_, origin_point_type_, member_, count_, sort_, false, GetUnitConversion(), + &geo_points); + + if (!s.ok()) { + return {Status::RedisExecErr, s.ToString()}; + } + *output = generateOutput(geo_points); + + return Status::OK(); + } + + protected: + double radius_ = 0; + double height_ = 0; + double width_ = 0; + int count_ = 0; + double longitude_ = 0; + double latitude_ = 0; + std::string member_; + std::string key_; + DistanceSort sort_ = kSortNone; + GeoShapeType shape_type_ = kGeoShapeTypeNone; + OriginPointType origin_point_type_ = kNone; + GeoShape geo_shape_; + + Status setShapeType(GeoShapeType shape_type) { + if (shape_type_ != kGeoShapeTypeNone) { + return {Status::RedisParseErr, "please use only one of BYBOX or BYRADIUS"}; + } + shape_type_ = shape_type; + return Status::OK(); + } + + Status setOriginType(OriginPointType origin_point_type) { + if (origin_point_type_ != kNone) { + return {Status::RedisParseErr, "please use only one of FROMMEMBER or FROMLONLAT"}; + } + origin_point_type_ = origin_point_type; + return Status::OK(); + } + + Status createGeoShape() { + if (shape_type_ == kGeoShapeTypeNone) { + return {Status::RedisParseErr, "please use BYBOX or BYRADIUS"}; + } + geo_shape_.type = shape_type_; + geo_shape_.conversion = GetUnitConversion(); + + if (shape_type_ == kGeoShapeTypeCircular) { + geo_shape_.radius = radius_; + } else { + geo_shape_.width = width_; + geo_shape_.height = height_; + } + + if (origin_point_type_ == kLongLat) { + geo_shape_.xy[0] = longitude_; + geo_shape_.xy[1] = latitude_; + } + return Status::OK(); + } + + std::string generateOutput(const std::vector &geo_points) { + int result_length = static_cast(geo_points.size()); + int returned_items_count = (count_ == 0 || result_length < count_) ? result_length : count_; + std::vector output; + output.reserve(returned_items_count); + for (int i = 0; i < returned_items_count; i++) { + auto geo_point = geo_points[i]; + if (!with_coord_ && !with_hash_ && !with_dist_) { + output.emplace_back(redis::BulkString(geo_point.member)); + } else { + std::vector one; + one.emplace_back(redis::BulkString(geo_point.member)); + if (with_dist_) { + one.emplace_back(redis::BulkString(util::Float2String(GetDistanceByUnit(geo_point.dist)))); + } + if (with_hash_) { + one.emplace_back(redis::BulkString(util::Float2String(geo_point.score))); + } + if (with_coord_) { + one.emplace_back(redis::MultiBulkString( + {util::Float2String(geo_point.longitude), util::Float2String(geo_point.latitude)})); + } + output.emplace_back(redis::Array(one)); + } + } + return redis::Array(output); + } + + private: + bool with_coord_ = false; + bool with_dist_ = false; + bool with_hash_ = false; +}; + +class CommandGeoSearchStore : public CommandGeoSearch { + public: + Status Parse(const std::vector &args) override { + CommandParser parser(args, 1); + store_key_ = GET_OR_RET(parser.TakeStr()); + key_ = GET_OR_RET(parser.TakeStr()); + + while (parser.Good()) { + if (parser.EatEqICase("frommember")) { + auto s = setOriginType(kMember); + if (!s.IsOK()) return s; + member_ = GET_OR_RET(parser.TakeStr()); + } else if (parser.EatEqICase("fromlonlat")) { + auto s = setOriginType(kLongLat); + if (!s.IsOK()) return s; + + longitude_ = GET_OR_RET(parser.TakeFloat()); + latitude_ = GET_OR_RET(parser.TakeFloat()); + s = ValidateLongLat(&longitude_, &latitude_); + if (!s.IsOK()) return s; + } else if (parser.EatEqICase("byradius")) { + auto s = setShapeType(kGeoShapeTypeCircular); + if (!s.IsOK()) return s; + radius_ = GET_OR_RET(parser.TakeFloat()); + std::string distance_raw = GET_OR_RET(parser.TakeStr()); + s = ParseDistanceUnit(distance_raw); + if (!s.IsOK()) return s; + } else if (parser.EatEqICase("bybox")) { + auto s = setShapeType(kGeoShapeTypeRectangular); + if (!s.IsOK()) return s; + width_ = GET_OR_RET(parser.TakeFloat()); + height_ = GET_OR_RET(parser.TakeFloat()); + std::string distance_raw = GET_OR_RET(parser.TakeStr()); + s = ParseDistanceUnit(distance_raw); + if (!s.IsOK()) return s; + } else if (parser.EatEqICase("asc") && sort_ == kSortNone) { + sort_ = kSortASC; + } else if (parser.EatEqICase("desc") && sort_ == kSortNone) { + sort_ = kSortDESC; + } else if (parser.EatEqICase("count")) { + count_ = GET_OR_RET(parser.TakeInt(NumericRange{1, std::numeric_limits::max()})); + } else if (parser.EatEqICase("storedist")) { + store_distance_ = true; + } else { + return {Status::RedisParseErr, "Invalid argument given"}; + } + } + + if (member_ != "" && longitude_ != 0 && latitude_ != 0) { + return {Status::RedisParseErr, "please use only one of FROMMEMBER or FROMLONLAT"}; + } + + auto s = createGeoShape(); + if (!s.IsOK()) { + return s; + } + return Commander::Parse(args); + } + + Status Execute(Server *svr, Connection *conn, std::string *output) override { + std::vector geo_points; + redis::Geo geo_db(svr->storage, conn->GetNamespace()); + + auto s = geo_db.SearchStore(args_[2], geo_shape_, origin_point_type_, member_, count_, sort_, store_key_, + store_distance_, GetUnitConversion(), &geo_points); + + if (!s.ok()) { + return {Status::RedisExecErr, s.ToString()}; + } + *output = redis::Integer(geo_points.size()); + return Status::OK(); + } + + private: + bool store_distance_ = false; + std::string store_key_; +}; + class CommandGeoRadiusByMember : public CommandGeoRadius { public: CommandGeoRadiusByMember() = default; @@ -406,7 +655,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr("geoadd", -5, "write", 1, 1, MakeCmdAttr("georadius", -6, "write", 1, 1, 1), MakeCmdAttr("georadiusbymember", -5, "write", 1, 1, 1), MakeCmdAttr("georadius_ro", -6, "read-only", 1, 1, 1), - MakeCmdAttr("georadiusbymember_ro", -5, "read-only", 1, 1, - 1), ) + MakeCmdAttr("georadiusbymember_ro", -5, "read-only", 1, 1, 1), + MakeCmdAttr("geosearch", -7, "read-only", 1, 1, 1), + MakeCmdAttr("geosearchstore", -8, "write", 1, 1, 1)) } // namespace redis diff --git a/src/types/geohash.cc b/src/types/geohash.cc index 3199df76675..120868f619b 100644 --- a/src/types/geohash.cc +++ b/src/types/geohash.cc @@ -329,47 +329,49 @@ uint8_t GeoHashHelper::EstimateStepsByRadius(double range_meters, double lat) { /* Return the bounding box of the search area centered at latitude,longitude * having a radius of radius_meter. bounds[0] - bounds[2] is the minimum * and maxium longitude, while bounds[1] - bounds[3] is the minimum and - * maximum latitude. - * - * This function does not behave correctly with very large radius values, for - * instance for the coordinates 81.634948934258375 30.561509253718668 and a - * radius of 7083 kilometers, it reports as bounding boxes: - * - * min_lon 7.680495, min_lat -33.119473, max_lon 155.589402, max_lat 94.242491 - * - * However, for instance, a min_lon of 7.680495 is not correct, because the - * point -1.27579540014266968 61.33421815228281559 is at less than 7000 - * kilometers away. - * - * Since this function is currently only used as an optimization, the - * optimization is not used for very big radiuses, however the function - * should be fixed. */ -int GeoHashHelper::BoundingBox(double longitude, double latitude, double radius_meters, double *bounds) { - if (!bounds) return 0; - - bounds[0] = longitude - RadDeg(radius_meters / EARTH_RADIUS_IN_METERS / cos(DegRad(latitude))); - bounds[2] = longitude + RadDeg(radius_meters / EARTH_RADIUS_IN_METERS / cos(DegRad(latitude))); - bounds[1] = latitude - RadDeg(radius_meters / EARTH_RADIUS_IN_METERS); - bounds[3] = latitude + RadDeg(radius_meters / EARTH_RADIUS_IN_METERS); + * maximum latitude. */ +int GeoHashHelper::BoundingBox(GeoShape *geo_shape) { + if (!geo_shape->bounds) return 0; + double longitude = geo_shape->xy[0]; + double latitude = geo_shape->xy[1]; + double height = + geo_shape->conversion * (geo_shape->type == kGeoShapeTypeCircular ? geo_shape->radius : geo_shape->height / 2); + double width = + geo_shape->conversion * (geo_shape->type == kGeoShapeTypeCircular ? geo_shape->radius : geo_shape->width / 2); + + const double lat_delta = RadDeg(height / EARTH_RADIUS_IN_METERS); + const double long_delta_top = RadDeg(width / EARTH_RADIUS_IN_METERS / cos(DegRad(latitude + lat_delta))); + const double long_delta_bottom = RadDeg(width / EARTH_RADIUS_IN_METERS / cos(DegRad(latitude - lat_delta))); + + bool is_in_southern_hemisphere = latitude < 0; + geo_shape->bounds[0] = is_in_southern_hemisphere ? longitude - long_delta_bottom : longitude - long_delta_top; + geo_shape->bounds[2] = is_in_southern_hemisphere ? longitude + long_delta_bottom : longitude + long_delta_top; + geo_shape->bounds[1] = latitude - lat_delta; + geo_shape->bounds[3] = latitude + lat_delta; return 1; } -/* Return a set of areas (center + 8) that are able to cover a range query - * for the specified position and radius. */ -GeoHashRadius GeoHashHelper::GetAreasByRadius(double longitude, double latitude, double radius_meters) { +GeoHashRadius GeoHashHelper::GetAreasByShapeWGS84(GeoShape &geo_shape) { GeoHashRange long_range, lat_range; GeoHashRadius radius; GeoHashBits hash; GeoHashNeighbors neighbors; GeoHashArea area; double min_lon = NAN, max_lon = NAN, min_lat = NAN, max_lat = NAN; - double bounds[4]; - BoundingBox(longitude, latitude, radius_meters, bounds); - min_lon = bounds[0]; - min_lat = bounds[1]; - max_lon = bounds[2]; - max_lat = bounds[3]; + BoundingBox(&geo_shape); + min_lon = geo_shape.bounds[0]; + min_lat = geo_shape.bounds[1]; + max_lon = geo_shape.bounds[2]; + max_lat = geo_shape.bounds[3]; + + double longitude = geo_shape.xy[0]; + double latitude = geo_shape.xy[1]; + + double radius_meters = + geo_shape.conversion * (geo_shape.type == kGeoShapeTypeCircular + ? geo_shape.radius + : sqrt(pow((geo_shape.width / 2), 2) + pow((geo_shape.height / 2), 2))); int steps = EstimateStepsByRadius(radius_meters, latitude); @@ -434,10 +436,6 @@ GeoHashRadius GeoHashHelper::GetAreasByRadius(double longitude, double latitude, return radius; } -GeoHashRadius GeoHashHelper::GetAreasByRadiusWGS84(double longitude, double latitude, double radius_meters) { - return GetAreasByRadius(longitude, latitude, radius_meters); -} - GeoHashFix52Bits GeoHashHelper::Align52Bits(const GeoHashBits &hash) { uint64_t bits = hash.bits; bits <<= (52 - hash.step * 2); @@ -462,7 +460,19 @@ int GeoHashHelper::GetDistanceIfInRadius(double x1, double y1, double x2, double return 1; } +int GeoHashHelper::GetDistanceIfInBox(const double *bounds, double x1, double y1, double x2, double y2, + double *distance) { + if (x2 < bounds[0] || x2 > bounds[2] || y2 < bounds[1] || y2 > bounds[3]) return 0; + *distance = GetDistance(x1, y1, x2, y2); + return 1; +} + int GeoHashHelper::GetDistanceIfInRadiusWGS84(double x1, double y1, double x2, double y2, double radius, double *distance) { return GetDistanceIfInRadius(x1, y1, x2, y2, radius, distance); } + +int GeoHashHelper::GetDistanceIfInBoxWGS84(const double *bounds, double x1, double y1, double x2, double y2, + double *distance) { + return GetDistanceIfInBox(bounds, x1, y1, x2, y2, distance); +} diff --git a/src/types/geohash.h b/src/types/geohash.h index 092eca198d1..c15df783a33 100644 --- a/src/types/geohash.h +++ b/src/types/geohash.h @@ -59,6 +59,8 @@ enum GeoDirection { GEOHASH_NORT_EAST }; +enum GeoShapeType { kGeoShapeTypeNone = 0, kGeoShapeTypeCircular, kGeoShapeTypeRectangular }; + struct GeoHashBits { uint64_t bits = 0; uint8_t step = 0; @@ -94,6 +96,16 @@ struct GeoHashRadius { GeoHashNeighbors neighbors; }; +struct GeoShape { + GeoShapeType type; + double xy[2]; + double conversion; + double bounds[4]; + double radius; + double height; + double width; +}; + inline constexpr bool HASHISZERO(const GeoHashBits &r) { return !r.bits && !r.step; } inline constexpr bool RANGEISZERO(const GeoHashRange &r) { return !bool(r.max) && !bool(r.min); } inline constexpr bool RANGEPISZERO(const GeoHashRange *r) { return !r || RANGEISZERO(*r); } @@ -122,11 +134,13 @@ void GeohashNeighbors(const GeoHashBits *hash, GeoHashNeighbors *neighbors); class GeoHashHelper { public: static uint8_t EstimateStepsByRadius(double range_meters, double lat); - static int BoundingBox(double longitude, double latitude, double radius_meters, double *bounds); - static GeoHashRadius GetAreasByRadius(double longitude, double latitude, double radius_meters); - static GeoHashRadius GetAreasByRadiusWGS84(double longitude, double latitude, double radius_meters); + static int BoundingBox(GeoShape *geo_shape); + static GeoHashRadius GetAreasByShapeWGS84(GeoShape &geo_shape); static GeoHashFix52Bits Align52Bits(const GeoHashBits &hash); static double GetDistance(double lon1d, double lat1d, double lon2d, double lat2d); static int GetDistanceIfInRadius(double x1, double y1, double x2, double y2, double radius, double *distance); + static int GetDistanceIfInBox(const double *bounds, double x1, double y1, double x2, double y2, double *distance); static int GetDistanceIfInRadiusWGS84(double x1, double y1, double x2, double y2, double radius, double *distance); + static int GetDistanceIfInBoxWGS84(const double *bounds, double x1, double y1, double x2, double y2, + double *distance); }; diff --git a/src/types/redis_geo.cc b/src/types/redis_geo.cc index bf179ba3b5d..e7c343c5d9b 100644 --- a/src/types/redis_geo.cc +++ b/src/types/redis_geo.cc @@ -79,30 +79,73 @@ rocksdb::Status Geo::Pos(const Slice &user_key, const std::vector &member rocksdb::Status Geo::Radius(const Slice &user_key, double longitude, double latitude, double radius_meters, int count, DistanceSort sort, const std::string &store_key, bool store_distance, double unit_conversion, std::vector *geo_points) { + GeoShape geo_shape; + geo_shape.type = kGeoShapeTypeCircular; + geo_shape.xy[0] = longitude; + geo_shape.xy[1] = latitude; + geo_shape.radius = radius_meters; + geo_shape.conversion = 1; + + std::string dummy_member; + return SearchStore(user_key, geo_shape, kLongLat, dummy_member, count, sort, store_key, store_distance, + unit_conversion, geo_points); +} + +rocksdb::Status Geo::RadiusByMember(const Slice &user_key, const Slice &member, double radius_meters, int count, + DistanceSort sort, const std::string &store_key, bool store_distance, + double unit_conversion, std::vector *geo_points) { + GeoPoint geo_point; + auto s = Get(user_key, member, &geo_point); + if (!s.ok()) return s; + + return Radius(user_key, geo_point.longitude, geo_point.latitude, radius_meters, count, sort, store_key, + store_distance, unit_conversion, geo_points); +} + +rocksdb::Status Geo::Search(const Slice &user_key, GeoShape geo_shape, OriginPointType point_type, std::string &member, + int count, DistanceSort sort, bool store_distance, double unit_conversion, + std::vector *geo_points) { + return SearchStore(user_key, geo_shape, point_type, member, count, sort, "", store_distance, unit_conversion, + geo_points); +} + +rocksdb::Status Geo::SearchStore(const Slice &user_key, GeoShape geo_shape, OriginPointType point_type, + std::string &member, int count, DistanceSort sort, const std::string &store_key, + bool store_distance, double unit_conversion, std::vector *geo_points) { + if (point_type == kMember) { + GeoPoint geo_point; + auto s = Get(user_key, member, &geo_point); + if (!s.ok()) return s; + + geo_shape.xy[0] = geo_point.longitude; + geo_shape.xy[1] = geo_point.latitude; + } + std::string ns_key; AppendNamespacePrefix(user_key, &ns_key); ZSetMetadata metadata(false); rocksdb::Status s = ZSet::GetMetadata(ns_key, &metadata); if (!s.ok()) return s.IsNotFound() ? rocksdb::Status::OK() : s; - /* Get all neighbor geohash boxes for our radius search */ - GeoHashRadius georadius = GeoHashHelper::GetAreasByRadiusWGS84(longitude, latitude, radius_meters); + // Get neighbor geohash boxes for radius search + GeoHashRadius georadius = GeoHashHelper::GetAreasByShapeWGS84(geo_shape); - /* Search the zset for all matching points */ - membersOfAllNeighbors(user_key, georadius, longitude, latitude, radius_meters, geo_points); + // Get zset for all matching points + membersOfAllNeighbors(user_key, georadius, geo_shape, geo_points); - /* If no matching results, the user gets an empty reply. */ - if (geo_points->empty() && store_key.empty()) { + // if no matching results, give empty reply + if (geo_points->empty()) { return rocksdb::Status::OK(); } - /* Process [optional] requested sorting */ + // process [optional] sorting if (sort == kSortASC) { std::sort(geo_points->begin(), geo_points->end(), sortGeoPointASC); } else if (sort == kSortDESC) { std::sort(geo_points->begin(), geo_points->end(), sortGeoPointDESC); } + // storing if (!store_key.empty()) { auto result_length = static_cast(geo_points->size()); int64_t returned_items_count = (count == 0 || result_length < count) ? result_length : count; @@ -121,21 +164,9 @@ rocksdb::Status Geo::Radius(const Slice &user_key, double longitude, double lati ZSet::Add(store_key, ZAddFlags::Default(), &member_scores, &ret); } } - return rocksdb::Status::OK(); } -rocksdb::Status Geo::RadiusByMember(const Slice &user_key, const Slice &member, double radius_meters, int count, - DistanceSort sort, const std::string &store_key, bool store_distance, - double unit_conversion, std::vector *geo_points) { - GeoPoint geo_point; - auto s = Get(user_key, member, &geo_point); - if (!s.ok()) return s; - - return Radius(user_key, geo_point.longitude, geo_point.latitude, radius_meters, count, sort, store_key, - store_distance, unit_conversion, geo_points); -} - rocksdb::Status Geo::Get(const Slice &user_key, const Slice &member, GeoPoint *geo_point) { std::map geo_points; auto s = MGet(user_key, {member}, &geo_points); @@ -208,7 +239,7 @@ int Geo::decodeGeoHash(double bits, double *xy) { } /* Search all eight neighbors + self geohash box */ -int Geo::membersOfAllNeighbors(const Slice &user_key, GeoHashRadius n, double lon, double lat, double radius, +int Geo::membersOfAllNeighbors(const Slice &user_key, GeoHashRadius n, const GeoShape &geo_shape, std::vector *geo_points) { GeoHashBits neighbors[9]; unsigned int last_processed = 0; @@ -239,7 +270,7 @@ int Geo::membersOfAllNeighbors(const Slice &user_key, GeoHashRadius n, double lo neighbors[i].step == neighbors[last_processed].step) { continue; } - count += membersOfGeoHashBox(user_key, neighbors[i], geo_points, lon, lat, radius); + count += membersOfGeoHashBox(user_key, neighbors[i], geo_points, geo_shape); last_processed = i; } return count; @@ -248,12 +279,12 @@ int Geo::membersOfAllNeighbors(const Slice &user_key, GeoHashRadius n, double lo /* Obtain all members between the min/max of this geohash bounding box. * Populate a GeoArray of GeoPoints by calling getPointsInRange(). * Return the number of points added to the array. */ -int Geo::membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash, std::vector *geo_points, double lon, - double lat, double radius) { +int Geo::membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash, std::vector *geo_points, + const GeoShape &geo_shape) { GeoHashFix52Bits min = 0, max = 0; scoresOfGeoHashBox(hash, &min, &max); - return getPointsInRange(user_key, static_cast(min), static_cast(max), lon, lat, radius, geo_points); + return getPointsInRange(user_key, static_cast(min), static_cast(max), geo_shape, geo_points); } /* Compute the sorted set scores min (inclusive), max (exclusive) we should @@ -297,7 +328,7 @@ void Geo::scoresOfGeoHashBox(GeoHashBits hash, GeoHashFix52Bits *min, GeoHashFix * using multiple queries to the sorted set, that we later need to sort * via qsort. Similarly we need to be able to reject points outside the search * radius area ASAP in order to allocate and process more points than needed. */ -int Geo::getPointsInRange(const Slice &user_key, double min, double max, double lon, double lat, double radius, +int Geo::getPointsInRange(const Slice &user_key, double min, double max, const GeoShape &geo_shape, std::vector *geo_points) { /* include min in range; exclude max in range */ /* That's: min <= val < max */ @@ -311,7 +342,7 @@ int Geo::getPointsInRange(const Slice &user_key, double min, double max, double if (!s.ok()) return 0; for (const auto &member_score : member_scores) { - appendIfWithinRadius(geo_points, lon, lat, radius, member_score.score, member_score.member); + appendIfWithinShape(geo_points, geo_shape, member_score.score, member_score.member); } return 0; } @@ -344,6 +375,33 @@ bool Geo::appendIfWithinRadius(std::vector *geo_points, double lon, do return true; } +bool Geo::appendIfWithinShape(std::vector *geo_points, const GeoShape &geo_shape, double score, + const std::string &member) { + double distance = NAN, xy[2]; + if (!decodeGeoHash(score, xy)) return false; + if (geo_shape.type == kGeoShapeTypeCircular) { + if (!GeoHashHelper::GetDistanceIfInRadiusWGS84(geo_shape.xy[0], geo_shape.xy[1], xy[0], xy[1], + geo_shape.radius * geo_shape.conversion, &distance)) { + return false; + } + } else if (geo_shape.type == kGeoShapeTypeRectangular) { + if (!GeoHashHelper::GetDistanceIfInBoxWGS84(geo_shape.bounds, geo_shape.xy[0], geo_shape.xy[1], xy[0], xy[1], + &distance)) { + return false; + } + } + + /* Append the new element. */ + GeoPoint geo_point; + geo_point.longitude = xy[0]; + geo_point.latitude = xy[1]; + geo_point.dist = distance; + geo_point.member = member; + geo_point.score = score; + geo_points->emplace_back(geo_point); + return true; +} + bool Geo::sortGeoPointASC(const GeoPoint &gp1, const GeoPoint &gp2) { return gp1.dist < gp2.dist; } bool Geo::sortGeoPointDESC(const GeoPoint &gp1, const GeoPoint &gp2) { return gp1.dist >= gp2.dist; } diff --git a/src/types/redis_geo.h b/src/types/redis_geo.h index 9c01f02dae3..c73f580f8a9 100644 --- a/src/types/redis_geo.h +++ b/src/types/redis_geo.h @@ -43,6 +43,8 @@ enum DistanceSort { kSortDESC, }; +enum OriginPointType { kNone, kLongLat, kMember }; + // Structures represent points and array of points on the earth. struct GeoPoint { double longitude; @@ -68,7 +70,12 @@ class Geo : public ZSet { rocksdb::Status RadiusByMember(const Slice &user_key, const Slice &member, double radius_meters, int count, DistanceSort sort, const std::string &store_key, bool store_distance, double unit_conversion, std::vector *geo_points); - + rocksdb::Status Search(const Slice &user_key, GeoShape geo_shape, OriginPointType point_type, std::string &member, + int count, DistanceSort sort, bool store_distance, double unit_conversion, + std::vector *geo_points); + rocksdb::Status SearchStore(const Slice &user_key, GeoShape geo_shape, OriginPointType point_type, + std::string &member, int count, DistanceSort sort, const std::string &store_key, + bool store_distance, double unit_conversion, std::vector *geo_points); rocksdb::Status Get(const Slice &user_key, const Slice &member, GeoPoint *geo_point); rocksdb::Status MGet(const Slice &user_key, const std::vector &members, std::map *geo_points); @@ -76,16 +83,17 @@ class Geo : public ZSet { private: static int decodeGeoHash(double bits, double *xy); - int membersOfAllNeighbors(const Slice &user_key, GeoHashRadius n, double lon, double lat, double radius, + int membersOfAllNeighbors(const Slice &user_key, GeoHashRadius n, const GeoShape &geo_shape, std::vector *geo_points); - int membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash, std::vector *geo_points, double lon, - double lat, double radius); + int membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash, std::vector *geo_points, + const GeoShape &geo_shape); static void scoresOfGeoHashBox(GeoHashBits hash, GeoHashFix52Bits *min, GeoHashFix52Bits *max); - int getPointsInRange(const Slice &user_key, double min, double max, double lon, double lat, double radius, + int getPointsInRange(const Slice &user_key, double min, double max, const GeoShape &geo_shape, std::vector *geo_points); static bool appendIfWithinRadius(std::vector *geo_points, double lon, double lat, double radius, double score, const std::string &member); - + static bool appendIfWithinShape(std::vector *geo_points, const GeoShape &geo_shape, double score, + const std::string &member); static bool sortGeoPointASC(const GeoPoint &gp1, const GeoPoint &gp2); static bool sortGeoPointDESC(const GeoPoint &gp1, const GeoPoint &gp2); }; diff --git a/tests/gocase/unit/geo/geo_test.go b/tests/gocase/unit/geo/geo_test.go index c34dee8621b..fb447287470 100644 --- a/tests/gocase/unit/geo/geo_test.go +++ b/tests/gocase/unit/geo/geo_test.go @@ -139,6 +139,69 @@ func TestGeo(t *testing.T) { require.EqualValues(t, []interface{}{nil, nil, nil}, rdb.Do(ctx, "GEOHASH", "points", "a", "b", "c").Val()) }) + t.Run("GEOSEARCH simple", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.GeoAdd(ctx, "points", + &redis.GeoLocation{Name: "Washington", Longitude: -77.0369, Latitude: 38.9072}, + &redis.GeoLocation{Name: "Baltimore", Longitude: -76.6121893, Latitude: 39.2903848}, + &redis.GeoLocation{Name: "New York", Longitude: -74.0059413, Latitude: 40.7127837}).Err()) + require.EqualValues(t, []string([]string{"Washington", "Baltimore", "New York"}), + rdb.GeoSearch(ctx, "points", &redis.GeoSearchQuery{Radius: 500, RadiusUnit: "km", Member: "Washington"}).Val()) + }) + + t.Run("GEOSEARCH simple (desc sorted)", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.GeoAdd(ctx, "points", + &redis.GeoLocation{Name: "Washington", Longitude: -77.0369, Latitude: 38.9072}, + &redis.GeoLocation{Name: "Baltimore", Longitude: -76.6121893, Latitude: 39.2903848}, + &redis.GeoLocation{Name: "New York", Longitude: -74.0059413, Latitude: 40.7127837}).Err()) + require.EqualValues(t, []string([]string{"New York", "Baltimore", "Washington"}), + rdb.GeoSearch(ctx, "points", &redis.GeoSearchQuery{Radius: 500, RadiusUnit: "km", Member: "Washington", Sort: "DESC"}).Val()) + }) + + t.Run("GEOSEARCH with coordinates", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.GeoAdd(ctx, "points", + &redis.GeoLocation{Name: "Washington", Longitude: -77.0369, Latitude: 38.9072}, + &redis.GeoLocation{Name: "Baltimore", Longitude: -76.6121893, Latitude: 39.2903848}, + &redis.GeoLocation{Name: "New York", Longitude: -74.0059413, Latitude: 40.7127837}).Err()) + require.EqualValues(t, []string([]string{"Baltimore", "Washington"}), + rdb.GeoSearch(ctx, "points", &redis.GeoSearchQuery{Radius: 200, RadiusUnit: "km", Longitude: -77.0368707, Latitude: 38.9071923, Sort: "DESC"}).Val()) + }) + + t.Run("GEOSEARCH with BYBOX on LongLat", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.GeoAdd(ctx, "points", + &redis.GeoLocation{Name: "Washington", Longitude: -77.0369, Latitude: 38.9072}, + &redis.GeoLocation{Name: "Baltimore", Longitude: -76.6121893, Latitude: 39.2903848}, + &redis.GeoLocation{Name: "New York", Longitude: -74.0059413, Latitude: 40.7127837}, + &redis.GeoLocation{Name: "Philadelphia", Longitude: -75.16521960, Latitude: 39.95258288}).Err()) + require.EqualValues(t, []string([]string{"Baltimore", "Washington"}), + rdb.GeoSearch(ctx, "points", &redis.GeoSearchQuery{BoxWidth: 200, BoxHeight: 200, BoxUnit: "km", Longitude: -77.0368707, Latitude: 38.9071923, Sort: "DESC"}).Val()) + }) + + t.Run("GEOSEARCH with BYBOX on member", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.GeoAdd(ctx, "points", + &redis.GeoLocation{Name: "Washington", Longitude: -77.0369, Latitude: 38.9072}, + &redis.GeoLocation{Name: "Baltimore", Longitude: -76.6121893, Latitude: 39.2903848}, + &redis.GeoLocation{Name: "New York", Longitude: -74.0059413, Latitude: 40.7127837}, + &redis.GeoLocation{Name: "Philadelphia", Longitude: -75.16521960, Latitude: 39.95258288}).Err()) + require.EqualValues(t, []string([]string{"Baltimore", "Washington"}), + rdb.GeoSearch(ctx, "points", &redis.GeoSearchQuery{BoxWidth: 200, BoxHeight: 200, BoxUnit: "km", Member: "Washington", Sort: "DESC"}).Val()) + }) + + t.Run("GEOSEARCHSTORE with BYRADIUS", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.GeoAdd(ctx, "points", + &redis.GeoLocation{Name: "Washington", Longitude: -77.0369, Latitude: 38.9072}, + &redis.GeoLocation{Name: "Baltimore", Longitude: -76.6121893, Latitude: 39.2903848}, + &redis.GeoLocation{Name: "New York", Longitude: -74.0059413, Latitude: 40.7127837}, + &redis.GeoLocation{Name: "Philadelphia", Longitude: -75.16521960, Latitude: 39.95258288}).Err()) + require.EqualValues(t, 2, + rdb.GeoSearchStore(ctx, "points", "points2", &redis.GeoSearchStoreQuery{GeoSearchQuery: redis.GeoSearchQuery{BoxWidth: 200, BoxHeight: 200, BoxUnit: "km", Longitude: -77.0368707, Latitude: 38.9071923, Sort: "DESC"}, StoreDist: false}).Val()) + }) + t.Run("GEOHASH is able to return geohash strings", func(t *testing.T) { require.NoError(t, rdb.Del(ctx, "points").Err()) require.NoError(t, rdb.GeoAdd(ctx, "points", &redis.GeoLocation{Name: "test", Longitude: -5.6, Latitude: 42.6}).Err()) @@ -198,6 +261,7 @@ func TestGeo(t *testing.T) { t.Run("GEORANGE STORE option: plain usage", func(t *testing.T) { require.NoError(t, rdb.Del(ctx, "points").Err()) + require.NoError(t, rdb.Del(ctx, "points2").Err()) require.NoError(t, rdb.GeoAdd(ctx, "points", &redis.GeoLocation{Name: "Palermo", Longitude: 13.361389, Latitude: 38.115556}, &redis.GeoLocation{Name: "Catania", Longitude: 15.087269, Latitude: 37.502669}).Err()) rdb.GeoRadiusStore(ctx, "points", 13.361389, 38.115556, &redis.GeoRadiusQuery{Radius: 500, Unit: "km", Store: "points2"}) require.EqualValues(t, rdb.ZRange(ctx, "points", 0, -1).Val(), rdb.ZRange(ctx, "points2", 0, -1).Val())