Skip to content

Commit

Permalink
refactor: geo related tcl tests (OpenAtomFoundation#2753)
Browse files Browse the repository at this point in the history
* modify geo.tcl ci

* modify go_test

* modify default.conf

* modify code based on review
  • Loading branch information
saz97 authored Jun 26, 2024
1 parent 8730946 commit 0d71b24
Show file tree
Hide file tree
Showing 7 changed files with 875 additions and 44 deletions.
45 changes: 34 additions & 11 deletions src/pika_geo.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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())) {
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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]);
Expand All @@ -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) {
Expand All @@ -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;
Expand All @@ -354,10 +364,18 @@ 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) {
if (db->storage()->Del({range.storekey}) > 0){
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;
Expand Down Expand Up @@ -426,6 +444,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;
Expand Down Expand Up @@ -555,6 +574,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: Invalid 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};
Expand Down
67 changes: 38 additions & 29 deletions src/pika_geohash_helper.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -79,7 +78,6 @@ uint8_t geohashEstimateStepsByRadius(double range_meters, double lat) {
step--;
}
}

/* Frame to valid range. */
if (step < 1) {
step = 1;
Expand Down Expand Up @@ -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;
}

Expand All @@ -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
Expand All @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down
3 changes: 2 additions & 1 deletion src/storage/src/redis_zsets.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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]);
Expand Down Expand Up @@ -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;
}
Expand Down
2 changes: 1 addition & 1 deletion tests/assets/default.conf
Original file line number Diff line number Diff line change
Expand Up @@ -567,4 +567,4 @@ cache-lfu-decay-time: 1
# Warning: Ensure that the Settings of rename-command on the master and slave servers are consistent
#
# Example:
# rename-command : FLUSHDB 360flushdb
# rename-command : FLUSHDB 360flushdb
2 changes: 1 addition & 1 deletion tests/integration/geo_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ var _ = Describe("Geo Commands", func() {
Expect(res.Err()).NotTo(HaveOccurred())
Expect(res.Val()).To(HaveLen(2))

Expect(res.Val()).To(Equal([]interface{}{[]interface{}{"Palermo", "190.4424", []interface{}{"13.361389338970184", "38.115556395496299"}}, []interface{}{"Catania", "56.4413", []interface{}{"15.087267458438873", "37.50266842333162"}}}))
Expect(res.Val()).To(Equal([]interface{}{[]interface{}{"Catania", "56.4413", []interface{}{"15.087267458438873", "37.50266842333162"}}, []interface{}{"Palermo", "190.4424", []interface{}{"13.361389338970184", "38.115556395496299"}}}))

})

Expand Down
2 changes: 1 addition & 1 deletion tests/unit/geo.tcl
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ start_server {tags {"geo"}} {

test {GEORADIUS with COUNT} {
r georadius nyc -73.9798091 40.7598464 10 km COUNT 3
} {{wtc one} {union square} {central park n/q/r}}
} {{central park n/q/r} 4545 {union square}}

test {GEORADIUS with COUNT but missing integer argument} {
catch {r georadius nyc -73.9798091 40.7598464 10 km COUNT} e
Expand Down
Loading

0 comments on commit 0d71b24

Please sign in to comment.