Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(geo): add geosearchstore #4272

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(geo): add geosearchstore #4272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 34 additions & 11 deletions src/server/zset_family.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2976,21 +2976,29 @@ void GeoSearchStoreGeneric(Transaction* tx, SinkReplyBuilder* builder, const Geo
rb->SendLong(smvec.size());
}
}
} // namespace

void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) {
// parse arguments
string_view key = ArgS(args, 0);
GeoShape shape = {};
void GeoSearchGeneric(CmdArgList args, bool store, const CommandContext& cmd_cntx) {
string_view key;
GeoSearchOpts geo_ops;
string_view member;
if (store) {
geo_ops.store = GeoStoreType::kStoreHash;
geo_ops.store_key = ArgS(args, 0);
key = ArgS(args, 1);
} else {
key = ArgS(args, 0);
}

// FROMMEMBER or FROMLONLAT is set
bool from_set = false;
string_view member;

// BYRADIUS or BYBOX is set
bool by_set = false;
GeoShape shape = {};

auto* builder = cmd_cntx.rb;
for (size_t i = 1; i < args.size(); ++i) {
size_t next_arg = (store) ? 2 : 1;
for (size_t i = next_arg; i < args.size(); ++i) {
string cur_arg = absl::AsciiStrToUpper(ArgS(args, i));

if (cur_arg == "FROMMEMBER") {
Expand Down Expand Up @@ -3085,13 +3093,15 @@ void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) {
geo_ops.any = true;
i++;
}
} else if (cur_arg == "WITHCOORD") {
} else if (!store && cur_arg == "WITHCOORD") {
geo_ops.withcoord = true;
} else if (cur_arg == "WITHDIST") {
} else if (!store && cur_arg == "WITHDIST") {
geo_ops.withdist = true;
} else if (cur_arg == "WITHHASH")
} else if (!store && cur_arg == "WITHHASH") {
geo_ops.withhash = true;
else {
} else if (store && cur_arg == "STOREDIST") {
geo_ops.store = GeoStoreType::kStoreDist;
} else {
return builder->SendError(kSyntaxErr);
}
}
Expand All @@ -3108,6 +3118,16 @@ void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) {
GeoSearchStoreGeneric(cmd_cntx.tx, builder, shape, key, member, geo_ops);
}

} // namespace

void ZSetFamily::GeoSearch(CmdArgList args, const CommandContext& cmd_cntx) {
GeoSearchGeneric(args, false, cmd_cntx);
}

void ZSetFamily::GeoSearchStore(CmdArgList args, const CommandContext& cmd_cntx) {
GeoSearchGeneric(args, true, cmd_cntx);
}

void ZSetFamily::GeoRadiusByMember(CmdArgList args, const CommandContext& cmd_cntx) {
GeoShape shape = {};
GeoSearchOpts geo_ops;
Expand Down Expand Up @@ -3243,6 +3263,7 @@ constexpr uint32_t kGeoHash = READ | GEO | SLOW;
constexpr uint32_t kGeoPos = READ | GEO | SLOW;
constexpr uint32_t kGeoDist = READ | GEO | SLOW;
constexpr uint32_t kGeoSearch = READ | GEO | SLOW;
constexpr uint32_t kGeoSearchStore = WRITE | GEO | SLOW;
constexpr uint32_t kGeoRadiusByMember = WRITE | GEO | SLOW;
} // namespace acl

Expand Down Expand Up @@ -3296,6 +3317,8 @@ void ZSetFamily::Register(CommandRegistry* registry) {
<< CI{"GEOPOS", CO::FAST | CO::READONLY, -2, 1, 1, acl::kGeoPos}.HFUNC(GeoPos)
<< CI{"GEODIST", CO::READONLY, -4, 1, 1, acl::kGeoDist}.HFUNC(GeoDist)
<< CI{"GEOSEARCH", CO::READONLY, -4, 1, 1, acl::kGeoSearch}.HFUNC(GeoSearch)
<< CI{"GEOSEARCHSTORE", CO::WRITE | CO::DENYOOM, -5, 1, 2, acl::kGeoSearchStore}.HFUNC(
Copy link
Contributor Author

@andydunstall andydunstall Dec 7, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm slightly guessing here about what CO:: params to use, not sure when CO::DENYOOM is needed etc - so just copied ZRANGESTORE which seems similar

GeoSearchStore)
<< CI{"GEORADIUSBYMEMBER", CO::WRITE | CO::STORE_LAST_KEY, -4, 1, 1, acl::kGeoRadiusByMember}
.HFUNC(GeoRadiusByMember);
}
Expand Down
1 change: 1 addition & 0 deletions src/server/zset_family.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ class ZSetFamily {
static void GeoPos(CmdArgList args, const CommandContext& cmd_cntx);
static void GeoDist(CmdArgList args, const CommandContext& cmd_cntx);
static void GeoSearch(CmdArgList args, const CommandContext& cmd_cntx);
static void GeoSearchStore(CmdArgList args, const CommandContext& cmd_cntx);
static void GeoRadiusByMember(CmdArgList args, const CommandContext& cmd_cntx);
};

Expand Down
23 changes: 23 additions & 0 deletions src/server/zset_family_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1132,6 +1132,29 @@ TEST_F(ZSetFamilyTest, GeoSearch) {
RespArray(ElementsAre(DoubleArg(9.1427), DoubleArg(38.7369))))))));
}

TEST_F(ZSetFamilyTest, GeoSearchStore) {
EXPECT_EQ(10, CheckedInt({"GEOADD", "Europe", "13.4050", "52.5200", "Berlin", "3.7038",
"40.4168", "Madrid", "9.1427", "38.7369", "Lisbon", "2.3522",
"48.8566", "Paris", "16.3738", "48.2082", "Vienna", "4.8952",
"52.3702", "Amsterdam", "10.7522", "59.9139", "Oslo", "23.7275",
"37.9838", "Athens", "19.0402", "47.4979", "Budapest", "6.2603",
"53.3498", "Dublin"}));

EXPECT_EQ(2, CheckedInt({"GEOSEARCHSTORE", "key1", "Europe", "FROMLONLAT", "13.4050", "52.5200",
"BYRADIUS", "500", "KM"}));

auto resp = Run({"ZRANGE", "key1", "0", "-1", "WITHSCORES"});
EXPECT_THAT(resp,
RespArray(ElementsAre("Berlin", "3673983950397063", "Dublin", "3678981558208417")));

EXPECT_EQ(2, CheckedInt({"GEOSEARCHSTORE", "key2", "Europe", "FROMLONLAT", "13.4050", "52.5200",
"BYRADIUS", "500", "KM", "STOREDIST"}));

resp = Run({"ZRANGE", "key2", "0", "-1", "WITHSCORES"});
EXPECT_THAT(resp, RespArray(ElementsAre("Berlin", DoubleArg(0.00017343178521311378), "Dublin",
DoubleArg(487.5619030644293))));
}

TEST_F(ZSetFamilyTest, GeoRadiusByMember) {
EXPECT_EQ(10, CheckedInt({"geoadd", "Europe", "13.4050", "52.5200", "Berlin", "3.7038",
"40.4168", "Madrid", "9.1427", "38.7369", "Lisbon", "2.3522",
Expand Down
Loading