This is an automated email from the ASF dual-hosted git repository.
hulk pushed a commit to branch unstable
in repository https://gitbox.apache.org/repos/asf/kvrocks.git
The following commit(s) were added to refs/heads/unstable by this push:
new 99cb7097 Add support for GEOSEARCH and GEOSEARCHSTORE (#1533)
99cb7097 is described below
commit 99cb709773fe15dcf7d374adce412c8d9c0e1f1e
Author: Uddeshya Singh <[email protected]>
AuthorDate: Mon Jul 24 12:42:07 2023 +0530
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 62dab23c..0b732833 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<std::string> &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<int>(NumericRange<int>{1,
std::numeric_limits<int>::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<GeoPoint> 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<GeoPoint> &geo_points) {
+ int result_length = static_cast<int>(geo_points.size());
+ int returned_items_count = (count_ == 0 || result_length < count_) ?
result_length : count_;
+ std::vector<std::string> 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<std::string> 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<std::string> &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<int>(NumericRange<int>{1,
std::numeric_limits<int>::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<GeoPoint> 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<CommandGeoAdd>("geoadd", -5, "write", 1, 1,
MakeCmdAttr<CommandGeoRadius>("georadius", -6,
"write", 1, 1, 1),
MakeCmdAttr<CommandGeoRadiusByMember>("georadiusbymember", -5, "write", 1, 1,
1),
MakeCmdAttr<CommandGeoRadiusReadonly>("georadius_ro",
-6, "read-only", 1, 1, 1),
-
MakeCmdAttr<CommandGeoRadiusByMemberReadonly>("georadiusbymember_ro", -5,
"read-only", 1, 1,
- 1), )
+
MakeCmdAttr<CommandGeoRadiusByMemberReadonly>("georadiusbymember_ro", -5,
"read-only", 1, 1, 1),
+ MakeCmdAttr<CommandGeoSearch>("geosearch", -7,
"read-only", 1, 1, 1),
+ MakeCmdAttr<CommandGeoSearchStore>("geosearchstore",
-8, "write", 1, 1, 1))
} // namespace redis
diff --git a/src/types/geohash.cc b/src/types/geohash.cc
index 3199df76..120868f6 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 092eca19..c15df783 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 bf179ba3..e7c343c5 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<Slice> &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<GeoPoint>
*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<GeoPoint> *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<GeoPoint> *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<GeoPoint> *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<int64_t>(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<GeoPoint> *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<std::string, GeoPoint> 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<GeoPoint> *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<GeoPoint> *geo_points, double lon,
- double lat, double radius) {
+int Geo::membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash,
std::vector<GeoPoint> *geo_points,
+ const GeoShape &geo_shape) {
GeoHashFix52Bits min = 0, max = 0;
scoresOfGeoHashBox(hash, &min, &max);
- return getPointsInRange(user_key, static_cast<double>(min),
static_cast<double>(max), lon, lat, radius, geo_points);
+ return getPointsInRange(user_key, static_cast<double>(min),
static_cast<double>(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<GeoPoint> *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<GeoPoint>
*geo_points, double lon, do
return true;
}
+bool Geo::appendIfWithinShape(std::vector<GeoPoint> *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 9c01f02d..c73f580f 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<GeoPoint>
*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<GeoPoint> *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<GeoPoint> *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<Slice>
&members,
std::map<std::string, GeoPoint> *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<GeoPoint> *geo_points);
- int membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash,
std::vector<GeoPoint> *geo_points, double lon,
- double lat, double radius);
+ int membersOfGeoHashBox(const Slice &user_key, GeoHashBits hash,
std::vector<GeoPoint> *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<GeoPoint> *geo_points);
static bool appendIfWithinRadius(std::vector<GeoPoint> *geo_points, double
lon, double lat, double radius,
double score, const std::string &member);
-
+ static bool appendIfWithinShape(std::vector<GeoPoint> *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 c34dee86..fb447287 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())