Add support for GEORADIUS

[azuredream]

#1878

[azuredream]

The test case failed because zrange returned [] while it should be ["Madrid", "Lisbon"].
I used GDB to debug it. During GEORADIUS STORE, the result was found but didn't trigger OpAdd.
That's because only shard0 went through store_cb while shard->size() was 2.

    auto store_cb = [&](Transaction* t, EngineShard* shard) {
      if (shard->shard_id() == dest_shard) {      //0!=1
        ZParams zparams;
        zparams.override = true;
        add_result =
            OpAdd(t->GetOpArgs(shard), zparams, geo_ops.store_key, ScoredMemberSpan{smvec}).value();
      }
      return OpStatus::OK;
    };

GEORADIUSMEMBER has an identity test case at line# 1134 and it passed. GDB showed that both shard0 and shard1 went through store_cb, and data was written to shard1.

[romange]

What are the keys involved in this operation?

[azuredream]

What are the keys involved in this operation?

2 keys. From key was "Europe", writing to store_key. Please see the trimmed test case below.
{"GEORADIUSBYMEMBER", "Europe", "Madrid", "700", "KM", "STORE", "store_key"}
{"GEORADIUS", "Europe", "3.7038", "40.4168", "700", "KM", "STORE", "store_key"}

I added some prints in GeoSearchStoreGeneric().
Based on the test logs below, test case for GeoRadiusByMember was able to see both shard 1 and 0.
But test case for GeoRadius was not able to see shard 1.

Test Logs:

[ RUN      ] ZSetFamilyTest.GeoRadiusByMember

from key: Europe, from shard: 0
 member checking:0
match
 member checking:1
dest key:store_key ,dest shard:1
smvec:
 3.47177e+15 Madrid
 3.47312e+15 Lisbon
checking:0
checking:1
match
[       OK ] ZSetFamilyTest.GeoRadiusByMember (18 ms)
[ RUN      ] ZSetFamilyTest.GeoRadius

from key: Europe, from shard: 0
no member checking:0
match
dest key:store_key ,dest shard:1
smvec:
 3.47177e+15 Madrid
 3.47312e+15 Lisbon
checking:0
.../dragonfly/src/server/zset_family_test.cc:1191: Failure
Value of: resp
Expected: resp array (value: 16-byte object <E4-E0 06-6D D5-55 00-00 CE-E0 06-6D D5-55 00-00>)
  Actual: [] (of type facade::RespExpr), whose given property is Vec: []

Trimmed Test cases:

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",
                            "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({"GEORADIUSBYMEMBER", "Europe", "Madrid", "700", "KM", "STORE", "store_key"}));
  auto 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")));
}

TEST_F(ZSetFamilyTest, GeoRadius) {
  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({"GEORADIUS", "Europe", "3.7038", "40.4168", "700", "KM", "STORE",
                           "store_key"}));
  auto 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")));
}

void GeoSearchStoreGeneric(ConnectionContext* cntx, const GeoShape& shape_ref, string_view key,
                           string_view member, const GeoSearchOpts& geo_ops) {
  GeoShape* shape = &(const_cast<GeoShape&>(shape_ref));
  auto* rb = static_cast<RedisReplyBuilder*>(cntx->reply_builder());

  std::cout<<std::endl;
  ShardId from_shard = Shard(key, shard_set->size());
  std::cout<<"from key: "<<key<<", from shard: "<<from_shard<<std::endl;

  if (!member.empty()) {
    // get shape.xy from member
    OpResult<double> member_score;
    auto cb = [&](Transaction* t, EngineShard* shard) {
      std::cout<<" member checking:"<<shard->shard_id()<<std::endl;
      if (shard->shard_id() == from_shard) {
        std::cout<<"match"<<std::endl;
        member_score = OpScore(t->GetOpArgs(shard), key, member);
      }
      return OpStatus::OK;
    };
    cntx->transaction->Execute(std::move(cb), false);
    auto member_sts = member_score.status();
    if (member_sts != OpStatus::OK) {
      std::cout<<"xxx"<<std::endl;
      cntx->transaction->Conclude();
      switch (member_sts) {
        case OpStatus::WRONG_TYPE:
          return cntx->SendError(kWrongTypeErr);
        case OpStatus::KEY_NOTFOUND:
          return rb->StartArray(0);
        case OpStatus::MEMBER_NOTFOUND:
          return cntx->SendError(kMemberNotFound);
        default:
          return cntx->SendError(member_sts);
      }
    }
    ScoreToLongLat(*member_score, shape->xy);
  } else {
    // verify key is valid
    OpResult<void> result;
    auto cb = [&](Transaction* t, EngineShard* shard) {
      std::cout<<"no member checking:"<<shard->shard_id()<<std::endl;
      if (shard->shard_id() == from_shard) {
        std::cout<<"match"<<std::endl;
        result = OpKeyExisted(t->GetOpArgs(shard), key);
      }
      return OpStatus::OK;
    };
    cntx->transaction->Execute(std::move(cb), false);
    auto result_sts = result.status();
    if (result_sts != OpStatus::OK) {
      std::cout<<"!!!"<<std::endl;
      cntx->transaction->Conclude();
      switch (result_sts) {
        case OpStatus::WRONG_TYPE:
          return cntx->SendError(kWrongTypeErr);
        case OpStatus::KEY_NOTFOUND:
          return rb->StartArray(0);
        default:
          return cntx->SendError(result_sts);
      }
    }
  }
...
  } else {
    // case 3: write mode, !kNoStore
    DCHECK(geo_ops.store == GeoStoreType::kStoreDist || geo_ops.store == GeoStoreType::kStoreHash);
    ShardId dest_shard = Shard(geo_ops.store_key, shard_set->size());
    std::cout<<"dest key:"<<geo_ops.store_key<<" ,dest shard:"<<dest_shard<<std::endl;
    DVLOG(1) << "store shard:" << dest_shard << ", key " << geo_ops.store_key;
    OpResult<AddResult> add_result;
    vector<ScoredMemberView> smvec;
    for (const auto& p : ga) {
      if (geo_ops.store == GeoStoreType::kStoreDist) {
        smvec.emplace_back(p.dist / geo_ops.conversion, p.member);
      } else {
        DCHECK(geo_ops.store == GeoStoreType::kStoreHash);
        smvec.emplace_back(p.score, p.member);
      }
    }
    std::cout<<"smvec:"<<std::endl;
    for (auto& si: smvec) {
      std::cout<<" "<<si.first<<" "<<si.second<<std::endl;
    }

    auto store_cb = [&](Transaction* t, EngineShard* shard) {
      std::cout<<"checking:"<<shard->shard_id()<<std::endl;
      if (shard->shard_id() == dest_shard) {
        std::cout<<"match"<<std::endl;
        ZParams zparams;
        zparams.override = true;
        add_result =
            OpAdd(t->GetOpArgs(shard), zparams, geo_ops.store_key, ScoredMemberSpan{smvec}).value();
      }
      return OpStatus::OK;
    };
    cntx->transaction->Execute(std::move(store_cb), true);

    rb->SendLong(smvec.size());
  }
}

[dranikpg]

That's a lot of code! 🙂

Quoting Redis docs:

It can be replaced by GEOSEARCH and GEOSEARCHSTORE with the BYRADIUS and FROMMEMBER arguments when migrating or writing new code.

So they already hint that GeoRadius is just a specification of GeoSearch. Maybe we can unite them somehow? They share count, asc/desc, withcord, withdist, storedist... Not mentioning parsing lon/lat and units

[azuredream]

That's a lot of code! 🙂

Quoting Redis docs:

It can be replaced by GEOSEARCH and GEOSEARCHSTORE with the BYRADIUS and FROMMEMBER arguments when migrating or writing new code.

So they already hint that GeoRadius is just a specification of GeoSearch. Maybe we can unite them somehow? They share count, asc/desc, withcord, withdist, storedist... Not mentioning parsing lon/lat and units

Thanks for reviewing. Updated as you suggested.