From 3f7e42b099f8ba0361502da66bd7b9712de1559e Mon Sep 17 00:00:00 2001 From: zixuan zhao Date: Mon, 27 Nov 2023 03:01:22 -0800 Subject: [PATCH] Add store test case for GeoRadiusByMember (#2210) * Add store test case for GeoRadiusByMember;Parsing code for STORE and STOREDIST --------- Signed-off-by: azuredream --- src/server/command_registry.cc | 2 ++ src/server/command_registry.h | 1 + src/server/transaction.cc | 18 +++++++++++++++--- src/server/zset_family.cc | 30 ++++++++++++++++++++++++++++++ src/server/zset_family_test.cc | 25 ++++++++++++++++++++++++- 5 files changed, 72 insertions(+), 4 deletions(-) diff --git a/src/server/command_registry.cc b/src/server/command_registry.cc index d71e10996..775ca2b84 100644 --- a/src/server/command_registry.cc +++ b/src/server/command_registry.cc @@ -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: diff --git a/src/server/command_registry.h b/src/server/command_registry.h index abf88ea23..c605b8fd2 100644 --- a/src/server/command_registry.h +++ b/src/server/command_registry.h @@ -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. diff --git a/src/server/transaction.cc b/src/server/transaction.cc index d1e656489..ebf29b034 100644 --- a/src/server/transaction.cc +++ b/src/server/transaction.cc @@ -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 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 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; } diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc index 5b99a1bca..183755423 100755 --- a/src/server/zset_family.cc +++ b/src/server/zset_family.cc @@ -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; 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); } diff --git a/src/server/zset_family_test.cc b/src/server/zset_family_test.cc index 8bdd05a99..1e7571057 100644 --- a/src/server/zset_family_test.cc +++ b/src/server/zset_family_test.cc @@ -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