refactor: geo related tcl tests

src/pika_geo.cc

@@ -10,6 +10,7 @@
 #include "pstd/include/pstd_string.h"
 
 #include "include/pika_geohash_helper.h"
+#include "rocksdb/status.h"
 
 void GeoAddCmd::DoInitial() {
   if (!CheckArg(argv_.size())) {
@@ -59,7 +60,7 @@ void GeoAddCmd::Do() {
   rocksdb::Status s = db_->storage()->ZAdd(key_, score_members, &count);
   if (s.ok()) {
     res_.AppendInteger(count);
-  } else if (s_.IsInvalidArgument()) {
+  } else if (s.IsInvalidArgument()) {
     res_.SetRes(CmdRes::kMultiKey);
   } else {
     res_.SetRes(CmdRes::kErrOther, s.ToString());
@@ -103,7 +104,7 @@ void GeoPosCmd::Do() {
     } else if (s.IsNotFound()) {
       res_.AppendStringLen(-1);
       continue;
-    } else if (s_.IsInvalidArgument()) {
+    } else if (s.IsInvalidArgument()) {
       res_.SetRes(CmdRes::kMultiKey);
       continue;
     } else {
@@ -163,13 +164,14 @@ void GeoDistCmd::Do() {
   double first_xy[2];
   double second_xy[2];
   rocksdb::Status s = db_->storage()->ZScore(key_, first_pos_, &first_score);
+
   if (s.ok()) {
     GeoHashBits hash = {.bits = static_cast<uint64_t>(first_score), .step = GEO_STEP_MAX};
     geohashDecodeToLongLatWGS84(hash, first_xy);
   } else if (s.IsNotFound()) {
     res_.AppendStringLen(-1);
     return;
-  } else if (s_.IsInvalidArgument()) {
+  } else if (s.IsInvalidArgument()) {
     res_.SetRes(CmdRes::kMultiKey);
     return;
   } else {
@@ -241,7 +243,7 @@ void GeoHashCmd::Do() {
     } else if (s.IsNotFound()) {
       res_.AppendStringLen(-1);
       continue;
-    } else if (s_.IsInvalidArgument()) {
+    } else if (s.IsInvalidArgument()) {
       res_.SetRes(CmdRes::kMultiKey);
       continue;
     } else {
@@ -300,6 +302,7 @@ static void GetAllNeighbors(const std::shared_ptr<DB>& db, std::string& key, Geo
     if (HASHISZERO(neighbors[i])) {
       continue;
     }
+
     min = geohashAlign52Bits(neighbors[i]);
     neighbors[i].bits++;
     max = geohashAlign52Bits(neighbors[i]);
@@ -312,8 +315,13 @@ static void GetAllNeighbors(const std::shared_ptr<DB>& db, std::string& key, Geo
     std::vector<storage::ScoreMember> score_members;
     s = db->storage()->ZRangebyscore(key, static_cast<double>(min), static_cast<double>(max), true, true, &score_members);
     if (!s.ok() && !s.IsNotFound()) {
-      res.SetRes(CmdRes::kErrOther, s.ToString());
-      return;
+      if (s.IsInvalidArgument()) {
+        res.SetRes(CmdRes::kMultiKey);
+        return;
+      } else {
+        res.SetRes(CmdRes::kErrOther, s.ToString());
+        return;
+      }
     }
     // Insert into result only if the point is within the search area.
     for (auto & score_member : score_members) {
@@ -339,12 +347,14 @@ static void GetAllNeighbors(const std::shared_ptr<DB>& db, std::string& key, Geo
     count_limit = static_cast<int32_t>(result.size());
   }
   // If using sort option
-  if (range.sort == Asc) {
-    std::sort(result.begin(), result.end(), sort_distance_asc);
-  } else if (range.sort == Desc) {
-    std::sort(result.begin(), result.end(), sort_distance_desc);
+  if (range.sort != Unsort) {
+    if (range.sort == Asc) {
+      std::sort(result.begin(), result.end(), sort_distance_asc);
+    } else if (range.sort == Desc) {
+      std::sort(result.begin(), result.end(), sort_distance_desc);
+    }
   }
-
+  
   if (range.store || range.storedist) {
     // Target key, create a sorted set with the results.
     std::vector<storage::ScoreMember> score_members;
@@ -354,10 +364,17 @@ static void GetAllNeighbors(const std::shared_ptr<DB>& db, std::string& key, Geo
       score_members.push_back({score, result[i].member});
     }
     int32_t count = 0;
+    int32_t card = db->storage()->Exists({range.storekey});
+    if (card) {
+      db->storage()->Del({range.storekey});
+      db->cache()->Del({range.storekey});
+    }
     s = db->storage()->ZAdd(range.storekey, score_members, &count);
     if (!s.ok()) {
       res.SetRes(CmdRes::kErrOther, s.ToString());
       return;
+    } else {
+      s = db->cache()->ZAdd(range.storekey, score_members);
     }
     res.AppendInteger(count_limit);
     return;
@@ -426,6 +443,7 @@ void GeoRadiusCmd::DoInitial() {
     return;
   }
   size_t pos = 6;
+  range_.sort = Asc;
   while (pos < argv_.size()) {
     if (strcasecmp(argv_[pos].c_str(), "withdist") == 0) {
       range_.withdist = true;
@@ -555,6 +573,10 @@ void GeoRadiusByMemberCmd::DoInitial() {
 void GeoRadiusByMemberCmd::Do() {
   double score = 0.0;
   rocksdb::Status s = db_->storage()->ZScore(key_, range_.member, &score);
+  if (s.IsNotFound() && !s.ToString().compare("NotFound: Invaild member")) {
+    res_.SetRes(CmdRes::kErrOther, "could not decode requested zset member");
+    return;
+  }
   if (s.ok()) {
     double xy[2];
     GeoHashBits hash = {.bits = static_cast<uint64_t>(score), .step = GEO_STEP_MAX};

src/pika_geohash_helper.cc

@@ -38,7 +38,6 @@
 #include "include/pika_geohash_helper.h"
 // #include "debugmacro.h"
 #include <cmath>
-
 #define D_R (M_PI / 180.0)
 #define R_MAJOR 6378137.0
 #define R_MINOR 6356752.3142
@@ -79,7 +78,6 @@ uint8_t geohashEstimateStepsByRadius(double range_meters, double lat) {
       step--;
     }
   }
-
   /* Frame to valid range. */
   if (step < 1) {
     step = 1;
@@ -112,11 +110,19 @@ int geohashBoundingBox(double longitude, double latitude, double radius_meters,
   if (!bounds) {
     return 0;
   }
+  double height = radius_meters;
+  double width = radius_meters;
+
+  const double lat_delta = rad_deg(height/EARTH_RADIUS_IN_METERS);
+  const double long_delta_top = rad_deg(width/EARTH_RADIUS_IN_METERS/cos(deg_rad(latitude+lat_delta)));
+  const double long_delta_bottom = rad_deg(width/EARTH_RADIUS_IN_METERS/cos(deg_rad(latitude-lat_delta)));
+
+  int southern_hemisphere = latitude < 0 ? 1 : 0;
+  bounds[0] = southern_hemisphere ? longitude-long_delta_bottom : longitude-long_delta_top;
+  bounds[2] = southern_hemisphere ? longitude+long_delta_bottom : longitude+long_delta_top;
+  bounds[1] = latitude - lat_delta;
+  bounds[3] = latitude + lat_delta;
 
-  bounds[0] = longitude - rad_deg(radius_meters / EARTH_RADIUS_IN_METERS / cos(deg_rad(latitude)));
-  bounds[2] = longitude + rad_deg(radius_meters / EARTH_RADIUS_IN_METERS / cos(deg_rad(latitude)));
-  bounds[1] = latitude - rad_deg(radius_meters / EARTH_RADIUS_IN_METERS);
-  bounds[3] = latitude + rad_deg(radius_meters / EARTH_RADIUS_IN_METERS);
   return 1;
 }
 
@@ -141,14 +147,12 @@ GeoHashRadius geohashGetAreasByRadius(double longitude, double latitude, double
   min_lat = bounds[1];
   max_lon = bounds[2];
   max_lat = bounds[3];
-
   steps = geohashEstimateStepsByRadius(radius_meters, latitude);
-
+  
   geohashGetCoordRange(&long_range, &lat_range);
   geohashEncode(&long_range, &lat_range, longitude, latitude, steps, &hash);
   geohashNeighbors(&hash, &neighbors);
   geohashDecode(long_range, lat_range, hash, &area);
-
   /* Check if the step is enough at the limits of the covered area.
    * Sometimes when the search area is near an edge of the
    * area, the estimated step is not small enough, since one of the
@@ -166,20 +170,19 @@ GeoHashRadius geohashGetAreasByRadius(double longitude, double latitude, double
     geohashDecode(long_range, lat_range, neighbors.east, &east);
     geohashDecode(long_range, lat_range, neighbors.west, &west);
 
-    if (geohashGetDistance(longitude, latitude, longitude, north.latitude.max) < radius_meters) {
+    if (north.latitude.max < max_lat) {
       decrease_step = 1;
     }
-    if (geohashGetDistance(longitude, latitude, longitude, south.latitude.min) < radius_meters) {
+    if (south.latitude.min > min_lat) {
       decrease_step = 1;
     }
-    if (geohashGetDistance(longitude, latitude, east.longitude.max, latitude) < radius_meters) {
+    if (east.longitude.max < max_lon) {
       decrease_step = 1;
     }
-    if (geohashGetDistance(longitude, latitude, west.longitude.min, latitude) < radius_meters) {
+    if (west.longitude.min > min_lon) {
       decrease_step = 1;
     }
   }
-
   if (steps > 1 && (decrease_step != 0)) {
     steps--;
     geohashEncode(&long_range, &lat_range, longitude, latitude, steps, &hash);
@@ -225,22 +228,28 @@ GeoHashFix52Bits geohashAlign52Bits(const GeoHashBits& hash) {
   bits <<= (52 - hash.step * 2);
   return bits;
 }
-
-/* Calculate distance using haversin great circle distance formula. */
+/* Calculate distance using simplified haversine great circle distance formula.
+ * Given longitude diff is 0 the asin(sqrt(a)) on the haversine is asin(sin(abs(u))).
+ * arcsin(sin(x)) equal to x when x ∈[−𝜋/2,𝜋/2]. Given latitude is between [−𝜋/2,𝜋/2]
+ * we can simplify arcsin(sin(x)) to x.
+ */
+double geohashGetLatDistance(double lat1d, double lat2d) {
+    return EARTH_RADIUS_IN_METERS * fabs(deg_rad(lat2d) - deg_rad(lat1d));
+}
+/* Calculate distance using haversine great circle distance formula. */
 double geohashGetDistance(double lon1d, double lat1d, double lon2d, double lat2d) {
-  double lat1r;
-  double lon1r;
-  double lat2r;
-  double lon2r;
-  double u;
-  double v;
-  lat1r = deg_rad(lat1d);
-  lon1r = deg_rad(lon1d);
-  lat2r = deg_rad(lat2d);
-  lon2r = deg_rad(lon2d);
-  u = sin((lat2r - lat1r) / 2);
-  v = sin((lon2r - lon1r) / 2);
-  return 2.0 * EARTH_RADIUS_IN_METERS * asin(sqrt(u * u + cos(lat1r) * cos(lat2r) * v * v));
+    double lat1r, lon1r, lat2r, lon2r, u, v, a;
+    lon1r = deg_rad(lon1d);
+    lon2r = deg_rad(lon2d);
+    v = sin((lon2r - lon1r) / 2);
+    /* if v == 0 we can avoid doing expensive math when lons are practically the same */
+    if (v == 0.0)
+        return geohashGetLatDistance(lat1d, lat2d);
+    lat1r = deg_rad(lat1d);
+    lat2r = deg_rad(lat2d);
+    u = sin((lat2r - lat1r) / 2);
+    a = u * u + cos(lat1r) * cos(lat2r) * v * v;
+    return 2.0 * EARTH_RADIUS_IN_METERS * asin(sqrt(a));
 }
 
 int geohashGetDistanceIfInRadius(double x1, double y1, double x2, double y2, double radius, double* distance) {

src/storage/src/redis_zsets.cc

@@ -517,7 +517,6 @@ Status Redis::ZRange(const Slice& key, int32_t start, int32_t stop, std::vector<
       }
       int32_t cur_index = 0;
       ScoreMember score_member;
-
       ZSetsScoreKey zsets_score_key(key, version, std::numeric_limits<double>::lowest(), Slice());
       KeyStatisticsDurationGuard guard(this, DataType::kZSets, key.ToString());
       rocksdb::Iterator* iter = db_->NewIterator(read_options, handles_[kZsetsScoreCF]);
@@ -1187,6 +1186,8 @@ Status Redis::ZScore(const Slice& key, const Slice& member, double* score) {
         uint64_t tmp = DecodeFixed64(data_value.data());
         const void* ptr_tmp = reinterpret_cast<const void*>(&tmp);
         *score = *reinterpret_cast<const double*>(ptr_tmp);
+      } else if (s.IsNotFound()) {
+        return Status::NotFound("Invalid member");
       } else {
         return s;
       }