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())


Reply via email to