// Copyright (c) 2015-present, Qihoo, Inc. All rights reserved. // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. An additional grant // of patent rights can be found in the PATENTS file in the same directory. #include #include #include "slash_string.h" #include "nemo.h" #include "pika_geo.h" #include "pika_server.h" #include "pika_slot.h" #include "pika_geohash_helper.h" extern PikaServer *g_pika_server; void GeoAddCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) { if (!ptr_info->CheckArg(argv.size())) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoAdd); return; } size_t argc = argv.size(); if ((argc - 2) % 3 != 0) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoAdd); return; } key_ = argv[1]; pos_.clear(); size_t index = 2; for (; index < argc; index += 3) { struct GeoPoint point; double longitude, latitude; if (!slash::string2d(argv[index].data(), argv[index].size(), &longitude)) { res_.SetRes(CmdRes::kInvalidFloat); return; } if (!slash::string2d(argv[index+1].data(), argv[index+1].size(), &latitude)) { res_.SetRes(CmdRes::kInvalidFloat); return; } point.member = argv[index+2]; point.longitude = longitude; point.latitude = latitude; pos_.push_back(point); } return; } void GeoAddCmd::Do() { nemo::Status s; int64_t count = 0, ret; const std::shared_ptr db = g_pika_server->db(); std::vector::const_iterator iter = pos_.begin(); for (; iter != pos_.end(); iter++) { // Convert coordinates to geohash GeoHashBits hash; geohashEncodeWGS84(iter->longitude, iter->latitude, GEO_STEP_MAX, &hash); GeoHashFix52Bits bits = geohashAlign52Bits(hash); // Convert uint64 to double double score; std::string str_bits = std::to_string(bits); slash::string2d(str_bits.data(), str_bits.size(), &score); s = db->ZAdd(key_, score, iter->member, &ret); if (s.ok()) { count += ret; } else { res_.SetRes(CmdRes::kErrOther, s.ToString()); return; } } SlotKeyAdd("z", key_); res_.AppendInteger(count); return; } void GeoPosCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) { if (!ptr_info->CheckArg(argv.size())) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoPos); return; } key_ = argv[1]; member_.clear(); size_t pos = 2; while (pos < argv.size()) { member_.push_back(argv[pos++]); } } void GeoPosCmd::Do() { double score; res_.AppendArrayLen(member_.size()); for (auto v : member_) { nemo::Status s = g_pika_server->db()->ZScore(key_, v, &score); if (s.ok()) { double xy[2]; GeoHashBits hash = { .bits = (uint64_t)score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, xy); res_.AppendArrayLen(2); char longitude[32]; int64_t len = slash::d2string(longitude, sizeof(longitude), xy[0]); res_.AppendStringLen(len); res_.AppendContent(longitude); char latitude[32]; len = slash::d2string(latitude, sizeof(latitude), xy[1]); res_.AppendStringLen(len); res_.AppendContent(latitude); } else if (s.IsNotFound()) { res_.AppendStringLen(-1); continue; } else { res_.SetRes(CmdRes::kErrOther, s.ToString()); continue; } } } static double length_converter(double meters, const std::string & unit) { if (unit == "m") { return meters; } else if (unit == "km") { return meters / 1000; } else if (unit == "ft") { return meters / 0.3048; } else if (unit == "mi") { return meters / 1609.34; } else { return -1; } } static bool check_unit(const std::string & unit) { if (unit == "m" || unit == "km" || unit == "ft" || unit == "mi") { return true; } else { return false; } } void GeoDistCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) { if (!ptr_info->CheckArg(argv.size())) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoDist); return; } if (argv.size() < 4) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoDist); return; } else if (argv.size() > 5) { res_.SetRes(CmdRes::kSyntaxErr); return; } key_ = argv[1]; first_pos_ = argv[2]; second_pos_ = argv[3]; if (argv.size() == 5) { unit_ = argv[4]; } else { unit_ = "m"; } if (!check_unit(unit_)) { res_.SetRes(CmdRes::kErrOther, "unsupported unit provided. please use m, km, ft, mi"); return; } } void GeoDistCmd::Do() { double first_score, second_score, first_xy[2], second_xy[2]; nemo::Status s = g_pika_server->db()->ZScore(key_, first_pos_, &first_score); if (s.ok()) { GeoHashBits hash = { .bits = (uint64_t)first_score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, first_xy); } else if (s.IsNotFound()) { res_.AppendStringLen(-1); return; } else { res_.SetRes(CmdRes::kErrOther, s.ToString()); return; } s = g_pika_server->db()->ZScore(key_, second_pos_, &second_score); if (s.ok()) { GeoHashBits hash = { .bits = (uint64_t)second_score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, second_xy); } else if (s.IsNotFound()) { res_.AppendStringLen(-1); return; } else { res_.SetRes(CmdRes::kErrOther, s.ToString()); return; } double distance = geohashGetDistance(first_xy[0], first_xy[1], second_xy[0], second_xy[1]); distance = length_converter(distance, unit_); char buf[32]; int64_t len = slash::d2string(buf, sizeof(buf), distance); res_.AppendStringLen(len); res_.AppendContent(buf); } void GeoHashCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) { if (!ptr_info->CheckArg(argv.size())) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoHash); return; } key_ = argv[1]; member_.clear(); size_t pos = 2; while (pos < argv.size()) { member_.push_back(argv[pos++]); } } void GeoHashCmd::Do() { const char * geoalphabet= "0123456789bcdefghjkmnpqrstuvwxyz"; res_.AppendArrayLen(member_.size()); for (auto v : member_) { double score; nemo::Status s = g_pika_server->db()->ZScore(key_, v, &score); if (s.ok()) { double xy[2]; GeoHashBits hash = { .bits = (uint64_t)score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, xy); GeoHashRange r[2]; GeoHashBits encode_hash; r[0].min = -180; r[0].max = 180; r[1].min = -90; r[1].max = 90; geohashEncode(&r[0], &r[1], xy[0], xy[1], 26, &encode_hash); char buf[12]; int i; for (i = 0; i < 11; i++) { int idx = (encode_hash.bits >> (52-((i+1)*5))) & 0x1f; buf[i] = geoalphabet[idx]; } buf[11] = '\0'; res_.AppendStringLen(11); res_.AppendContent(buf); continue; } else if (s.IsNotFound()) { res_.AppendStringLen(-1); continue; } else { res_.SetRes(CmdRes::kErrOther, s.ToString()); continue; } } } static bool sort_distance_asc(const NeighborPoint & pos1, const NeighborPoint & pos2) { return pos1.distance < pos2.distance; } static bool sort_distance_desc(const NeighborPoint & pos1, const NeighborPoint & pos2) { return pos1.distance > pos2.distance; } static void GetAllNeighbors(std::string & key, GeoRange & range, CmdRes & res) { nemo::Status s; double longitude = range.longitude, latitude = range.latitude, distance = range.distance; int count_limit = 0; // Convert other units to meters if (range.unit == "m") { distance = distance; } else if (range.unit == "km") { distance = distance * 1000; } else if (range.unit == "ft") { distance = distance * 0.3048; } else if (range.unit == "mi") { distance = distance * 1609.34; } else { distance = -1; } // Search the zset for all matching points GeoHashRadius georadius = geohashGetAreasByRadiusWGS84(longitude, latitude, distance); GeoHashBits neighbors[9]; neighbors[0] = georadius.hash; neighbors[1] = georadius.neighbors.north; neighbors[2] = georadius.neighbors.south; neighbors[3] = georadius.neighbors.east; neighbors[4] = georadius.neighbors.west; neighbors[5] = georadius.neighbors.north_east; neighbors[6] = georadius.neighbors.north_west; neighbors[7] = georadius.neighbors.south_east; neighbors[8] = georadius.neighbors.south_west; // For each neighbor, get all the matching // members and add them to the potential result list. std::vector result; for (size_t i = 0; i < sizeof(neighbors) / sizeof(*neighbors); i++) { GeoHashFix52Bits min, max; if (HASHISZERO(neighbors[i])) continue; min = geohashAlign52Bits(neighbors[i]); neighbors[i].bits++; max = geohashAlign52Bits(neighbors[i]); std::vector sm_v; s = g_pika_server->db()->ZRangebyscore(key, (double)min, (double)max, sm_v, false, false); if (!s.ok()) { res.SetRes(CmdRes::kErrOther, s.ToString()); return; } // Insert into result only if the point is within the search area. for (size_t i = 0; i < sm_v.size(); ++i) { double xy[2], real_distance; GeoHashBits hash = { .bits = (uint64_t)sm_v[i].score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, xy); if(geohashGetDistanceIfInRadiusWGS84(longitude, latitude, xy[0], xy[1], distance, &real_distance)) { NeighborPoint item; item.member = sm_v[i].member; item.score = sm_v[i].score; item.distance = real_distance; result.push_back(item); } } } // If using the count opiton if (range.count) { count_limit = range.count_limit; } else { count_limit = 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); } // For each the result res.AppendArrayLen(count_limit); for (int i = 0; i < count_limit; ++i) { if (range.option_num != 0) { res.AppendArrayLen(range.option_num+1); } // Member res.AppendStringLen(result[i].member.size()); res.AppendContent(result[i].member); // If using withdist option if (range.withdist) { double xy[2]; GeoHashBits hash = { .bits = (uint64_t)result[i].score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, xy); double distance = geohashGetDistance(longitude, latitude, xy[0], xy[1]); distance = length_converter(distance, range.unit); char buf[32]; int64_t len = slash::d2string(buf, sizeof(buf), distance); res.AppendStringLen(len); res.AppendContent(buf); } // If using withcoord option if (range.withcoord) { res.AppendArrayLen(2); double xy[2]; GeoHashBits hash = { .bits = (uint64_t)result[i].score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, xy); char longitude[32]; int64_t len = slash::d2string(longitude, sizeof(longitude), xy[0]); res.AppendStringLen(len); res.AppendContent(longitude); char latitude[32]; len = slash::d2string(latitude, sizeof(latitude), xy[1]); res.AppendStringLen(len); res.AppendContent(latitude); } // If using withhash option if (range.withhash) { char buf[32]; int64_t len = slash::d2string(buf, sizeof(buf), result[i].score); res.AppendStringLen(len); res.AppendContent(buf); } } } void GeoRadiusCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) { if (!ptr_info->CheckArg(argv.size())) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoRadius); return; } key_ = argv[1]; slash::string2d(argv[2].data(), argv[2].size(), &range_.longitude); slash::string2d(argv[3].data(), argv[3].size(), &range_.latitude); slash::string2d(argv[4].data(), argv[4].size(), &range_.distance); range_.unit = argv[5]; if (!check_unit(range_.unit)) { res_.SetRes(CmdRes::kErrOther, "unsupported unit provided. please use m, km, ft, mi"); return; } size_t pos = 6; while (pos < argv.size()) { if (!strcasecmp(argv[pos].c_str(), "withdist")) { range_.withdist = true; range_.option_num++; } else if (!strcasecmp(argv[pos].c_str(), "withhash")) { range_.withhash = true; range_.option_num++; } else if (!strcasecmp(argv[pos].c_str(), "withcoord")) { range_.withcoord = true; range_.option_num++; } else if (!strcasecmp(argv[pos].c_str(), "count")) { range_.count = true; range_.count_limit = std::stoi(argv[++pos]); } else if (!strcasecmp(argv[pos].c_str(), "asc")) { range_.sort = Asc; } else if (!strcasecmp(argv[pos].c_str(), "desc")) { range_.sort = Desc; } else { res_.SetRes(CmdRes::kSyntaxErr); return; } pos++; } } void GeoRadiusCmd::Do() { GetAllNeighbors(key_, range_, this->res_); } void GeoRadiusByMemberCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) { if (!ptr_info->CheckArg(argv.size())) { res_.SetRes(CmdRes::kWrongNum, kCmdNameGeoRadius); return; } key_ = argv[1]; range_.member = argv[2]; slash::string2d(argv[3].data(), argv[3].size(), &range_.distance); range_.unit = argv[4]; if (!check_unit(range_.unit)) { res_.SetRes(CmdRes::kErrOther, "unsupported unit provided. please use m, km, ft, mi"); return; } size_t pos = 5; while (pos < argv.size()) { if (!strcasecmp(argv[pos].c_str(), "withdist")) { range_.withdist = true; range_.option_num++; } else if (!strcasecmp(argv[pos].c_str(), "withhash")) { range_.withhash = true; range_.option_num++; } else if (!strcasecmp(argv[pos].c_str(), "withcoord")) { range_.withcoord = true; range_.option_num++; } else if (!strcasecmp(argv[pos].c_str(), "count")) { range_.count = true; range_.count_limit = std::stoi(argv[++pos]); } else if (!strcasecmp(argv[pos].c_str(), "asc")) { range_.sort = Asc; } else if (!strcasecmp(argv[pos].c_str(), "desc")) { range_.sort = Desc; } else { res_.SetRes(CmdRes::kSyntaxErr); return; } pos++; } } void GeoRadiusByMemberCmd::Do() { nemo::Status s; double score; s = g_pika_server->db()->ZScore(key_, range_.member, &score); if (s.ok()) { double xy[2]; GeoHashBits hash = { .bits = (uint64_t)score, .step = GEO_STEP_MAX }; geohashDecodeToLongLatWGS84(hash, xy); range_.longitude = xy[0]; range_.latitude = xy[1]; } GetAllNeighbors(key_, range_, this->res_); }