mirrors/dragonfly
1
0
Fork 0
mirror of https://github.com/dragonflydb/dragonfly.git synced 2025-05-10 18:05:44 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Add store test case for GeoRadiusByMember (#2210)

Browse source 
* Add store test case for GeoRadiusByMember;Parsing code for STORE and STOREDIST

---------

Signed-off-by: azuredream <zhaozixuan67@gmail.com>
This commit is contained in:
zixuan zhao 2023-11-27 03:01:22 -08:00 • committed by GitHub
parent 92614477b7
commit 3f7e42b099
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 72 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/server/command_registry.cc
View file

@ -193,6 +193,8 @@ const char* OptName(CO::CommandOpt fl) {
return "interleaved-keys";
case GLOBAL_TRANS:
return "global-trans";
case STORE_LAST_KEY:
return "store-last-key";
case VARIADIC_KEYS:
return "variadic-keys";
case NO_AUTOJOURNAL:

1
src/server/command_registry.h
View file

@ -36,6 +36,7 @@ enum CommandOpt : uint32_t {
HIDDEN = 1U << 10, // does not show in COMMAND command output
INTERLEAVED_KEYS = 1U << 11, // keys are interleaved with arguments
GLOBAL_TRANS = 1U << 12,
STORE_LAST_KEY = 1U << 13, // The command my have a store key as the last argument.
NO_AUTOJOURNAL = 1U << 15, // Skip automatically logging command to journal inside transaction.

18
src/server/transaction.cc
View file

@ -241,7 +241,6 @@ void Transaction::InitByKeys(KeyIndex key_index) {
bool is_stub = multi_ && multi_->role == SQUASHED_STUB;
if ((key_index.HasSingleKey() && !IsAtomicMulti()) || is_stub) {
DCHECK_GT(key_index.step, 0u);
// We don't have to split the arguments by shards, so we can copy them directly.
StoreKeysInArgs(key_index, needs_reverse_mapping);
@ -260,8 +259,8 @@ void Transaction::InitByKeys(KeyIndex key_index) {
}
shard_data_.resize(shard_set->size()); // shard_data isn't sparse, so we must allocate for all :(
CHECK(key_index.step == 1 || key_index.step == 2);
DCHECK(key_index.step == 1 || (args.size() % 2) == 0);
DCHECK(key_index.step == 1 || key_index.step == 2);
DCHECK(key_index.step != 2 || (args.size() % 2) == 0);
// Safe, because flow below is not preemptive.
auto& shard_index = tmp_space.GetShardIndex(shard_data_.size());
@ -1494,6 +1493,7 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
if (cid->first_key_pos() > 0) {
key_index.start = cid->first_key_pos() - 1;
int last = cid->last_key_pos();
if (num_custom_keys >= 0) {
key_index.end = key_index.start + num_custom_keys;
} else {
@ -1501,6 +1501,18 @@ OpResult<KeyIndex> DetermineKeys(const CommandId* cid, CmdArgList args) {
}
key_index.step = cid->opt_mask() & CO::INTERLEAVED_KEYS ? 2 : 1;
if (cid->opt_mask() & CO::STORE_LAST_KEY) {
string_view name{cid->name()};
if (name == "GEORADIUSBYMEMBER" && args.size() >= 5) {
// key member radius .. STORE destkey
string_view opt = ArgS(args, args.size() - 2);
if (absl::EqualsIgnoreCase(opt, "STORE") || absl::EqualsIgnoreCase(opt, "STOREDIST")) {
key_index.bonus = args.size() - 1;
}
}
}
return key_index;
}

30
src/server/zset_family.cc
View file

@ -45,6 +45,8 @@ static const char kFromMemberLonglatErr[] =
static const char kByRadiusBoxErr[] =
"BYRADIUS and BYBOX options at the same time are not compatible";
static const char kAscDescErr[] = "ASC and DESC options at the same time are not compatible";
static const char kStoreTypeErr[] =
"STORE and STOREDIST options at the same time are not compatible";
static const char kScoreNaN[] = "resulting score is not a number (NaN)";
static const char kFloatRangeErr[] = "min or max is not a float";
static const char kLexRangeErr[] = "min or max not valid string range item";
@ -69,6 +71,7 @@ struct GeoPoint {
using GeoArray = std::vector<GeoPoint>;
enum class Sorting { kUnsorted, kAsc, kDesc };
enum class GeoStoreType { kNoStore, kStoreHash, kStoreDist };
struct GeoSearchOpts {
double conversion = 0;
uint64_t count = 0;
@ -77,6 +80,7 @@ struct GeoSearchOpts {
bool withdist = 0;
bool withcoord = 0;
bool withhash = 0;
GeoStoreType store = GeoStoreType::kNoStore;
};
inline zrangespec GetZrangeSpec(bool reverse, const ZSetFamily::ScoreInterval& si) {
@ -2789,6 +2793,9 @@ bool MembersOfAllNeighbors(ConnectionContext* cntx, string_view key, const GeoHa
if (result_arrays.status() == OpStatus::WRONG_TYPE) {
(*cntx)->SendError(kWrongTypeErr);
return false;
} else if (result_arrays.status() == OpStatus::KEY_NOTFOUND) {
(*cntx)->SendError("Member not found");
return false;
}
// filter potential result list
@ -3049,6 +3056,7 @@ void ZSetFamily::GeoRadiusByMember(CmdArgList args, ConnectionContext* cntx) {
return (*cntx)->SendError("unsupported unit provided. please use M, KM, FT, MI");
}
shape.type = CIRCULAR_TYPE;
string_view store_key;
for (size_t i = 4; i < args.size(); ++i) {
ToUpper(&args[i]);
@ -3083,6 +3091,28 @@ void ZSetFamily::GeoRadiusByMember(CmdArgList args, ConnectionContext* cntx) {
geo_ops.withdist = true;
} else if (cur_arg == "WITHHASH") {
geo_ops.withhash = true;
} else if (cur_arg == "STORE") {
if (geo_ops.store != GeoStoreType::kNoStore) {
return (*cntx)->SendError(kStoreTypeErr);
}
if (i + 1 < args.size()) {
store_key = ArgS(args, i + 1);
geo_ops.store = GeoStoreType::kStoreHash;
i++;
} else {
return (*cntx)->SendError(kSyntaxErr);
}
} else if (cur_arg == "STOREDIST") {
if (geo_ops.store != GeoStoreType::kNoStore) {
return (*cntx)->SendError(kStoreTypeErr);
}
if (i + 1 < args.size()) {
store_key = ArgS(args, i + 1);
geo_ops.store = GeoStoreType::kStoreDist;
i++;
} else {
return (*cntx)->SendError(kSyntaxErr);
}
} else {
return (*cntx)->SendError(kSyntaxErr);
}

25
src/server/zset_family_test.cc
View file

@ -862,6 +862,14 @@ TEST_F(ZSetFamilyTest, GeoSearch) {
RespArray(ElementsAre("Dublin", DoubleArg(487.5619030644293), "3678981558208417",
RespArray(ElementsAre(DoubleArg(6.2603), DoubleArg(53.3498))))))));
resp = Run({"GEOSEARCH", "America", "FROMMEMBER", "Madrid", "BYRADIUS", "700", "KM", "WITHCOORD",
"WITHDIST"});
EXPECT_THAT(resp, ErrArg("Member not found"));
resp = Run({"GEOSEARCH", "America", "FROMLONLAT", "13.4050", "52.5200", "BYBOX", "1000", "1000",
"KM", "WITHCOORD", "WITHDIST"});
EXPECT_THAT(resp, ErrArg("Member not found"));
resp = Run({"GEOSEARCH", "Europe", "FROMLONLAT", "13.4050", "52.5200", "BYBOX", "1000", "1000",
"KM", "WITHCOORD", "WITHDIST"});
EXPECT_THAT(
@ -912,7 +920,6 @@ TEST_F(ZSetFamilyTest, GeoRadiusByMember) {
"52.3702", "Amsterdam", "10.7522", "59.9139", "Oslo", "23.7275",
"37.9838", "Athens", "19.0402", "47.4979", "Budapest", "6.2603",
"53.3498", "Dublin"}));
auto resp = Run({"GEORADIUSBYMEMBER", "Europe", "Madrid", "700", "KM", "WITHCOORD", "WITHDIST"});
EXPECT_THAT(
resp,
@ -922,6 +929,22 @@ TEST_F(ZSetFamilyTest, GeoRadiusByMember) {
RespArray(
ElementsAre("Lisbon", "502.20769462704084",
RespArray(ElementsAre("9.142698347568512", "38.736900197448534")))))));
GTEST_SKIP();
EXPECT_EQ(
2, CheckedInt({"GEORADIUSBYMEMBER", "Europe", "Madrid", "700", "KM", "STORE", "store_key"}));
resp = Run({"ZRANGE", "store_key", "0", "-1"});
EXPECT_THAT(resp, RespArray(ElementsAre("Madrid", "Lisbon")));
resp = Run({"ZRANGE", "store_key", "0", "-1", "WITHSCORES"});
EXPECT_THAT(resp,
RespArray(ElementsAre("Madrid", "3471766229222696", "Lisbon", "3473121093062745")));
EXPECT_EQ(2, CheckedInt({"GEORADIUSBYMEMBER", "Europe", "Madrid", "700", "KM", "STOREDIST",
"store_dist_key"}));
resp = Run({"ZRANGE", "store_dist_key", "0", "-1", "WITHSCORES"});
EXPECT_THAT(resp, RespArray(ElementsAre("Madrid", "0", "Lisbon", "502.20769462704084")));
}
} // namespace dfly