GEOSEARCH - ANY option, for limited search that returns ASAP (#8259)

Support ANY option to return some results that match the criteria ASAP,
without a complete search and implicit sorting.
This commit is contained in:
Yang Bodong 2021-01-09 00:29:44 +08:00 committed by GitHub
parent 814aad65f1
commit ea5350c5ec
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 40 additions and 14 deletions

View File

@ -249,7 +249,7 @@ int geoAppendIfWithinShape(geoArray *ga, GeoShape *shape, double score, sds memb
* using multiple queries to the sorted set, that we later need to sort
* via qsort. Similarly we need to be able to reject points outside the search
* radius area ASAP in order to allocate and process more points than needed. */
int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geoArray *ga) {
int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geoArray *ga, unsigned long limit) {
/* minex 0 = include min in range; maxex 1 = exclude max in range */
/* That's: min <= val < max */
zrangespec range = { .min = min, .max = max, .minex = 0, .maxex = 1 };
@ -283,6 +283,7 @@ int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geo
sdsnewlen(vstr,vlen);
if (geoAppendIfWithinShape(ga,shape,score,member)
== C_ERR) sdsfree(member);
if (ga->used && limit && ga->used >= limit) break;
zzlNext(zl, &eptr, &sptr);
}
} else if (zobj->encoding == OBJ_ENCODING_SKIPLIST) {
@ -304,6 +305,7 @@ int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geo
ele = sdsdup(ele);
if (geoAppendIfWithinShape(ga,shape,ln->score,ele)
== C_ERR) sdsfree(ele);
if (ga->used && limit && ga->used >= limit) break;
ln = ln->level[0].forward;
}
}
@ -342,15 +344,15 @@ void scoresOfGeoHashBox(GeoHashBits hash, GeoHashFix52Bits *min, GeoHashFix52Bit
/* Obtain all members between the min/max of this geohash bounding box.
* Populate a geoArray of GeoPoints by calling geoGetPointsInRange().
* Return the number of points added to the array. */
int membersOfGeoHashBox(robj *zobj, GeoHashBits hash, geoArray *ga, GeoShape *shape) {
int membersOfGeoHashBox(robj *zobj, GeoHashBits hash, geoArray *ga, GeoShape *shape, unsigned long limit) {
GeoHashFix52Bits min, max;
scoresOfGeoHashBox(hash,&min,&max);
return geoGetPointsInRange(zobj, min, max, shape, ga);
return geoGetPointsInRange(zobj, min, max, shape, ga, limit);
}
/* Search all eight neighbors + self geohash box */
int membersOfAllNeighbors(robj *zobj, GeoHashRadius n, GeoShape *shape, geoArray *ga) {
int membersOfAllNeighbors(robj *zobj, GeoHashRadius n, GeoShape *shape, geoArray *ga, unsigned long limit) {
GeoHashBits neighbors[9];
unsigned int i, count = 0, last_processed = 0;
int debugmsg = 0;
@ -401,7 +403,8 @@ int membersOfAllNeighbors(robj *zobj, GeoHashRadius n, GeoShape *shape, geoArray
D("Skipping processing of %d, same as previous\n",i);
continue;
}
count += membersOfGeoHashBox(zobj, neighbors[i], ga, shape);
if (ga->used && limit && ga->used >= limit) break;
count += membersOfGeoHashBox(zobj, neighbors[i], ga, shape, limit);
last_processed = i;
}
return count;
@ -500,12 +503,12 @@ void geoaddCommand(client *c) {
#define GEOSEARCHSTORE (1<<4) /* GEOSEARCHSTORE just accept STOREDIST option */
/* GEORADIUS key x y radius unit [WITHDIST] [WITHHASH] [WITHCOORD] [ASC|DESC]
* [COUNT count] [STORE key] [STOREDIST key]
* [COUNT count [ANY]] [STORE key] [STOREDIST key]
* GEORADIUSBYMEMBER key member radius unit ... options ...
* GEOSEARCH key [FROMMEMBER member] [FROMLONLAT long lat] [BYRADIUS radius unit]
* [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT count] [ASC|DESC]
* [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT count [ANY]] [ASC|DESC]
* GEOSEARCHSTORE dest_key src_key [FROMMEMBER member] [FROMLONLAT long lat] [BYRADIUS radius unit]
* [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT count] [ASC|DESC] [STOREDIST]
* [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT count [ANY]] [ASC|DESC] [STOREDIST]
* */
void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
robj *storekey = NULL;
@ -550,7 +553,8 @@ void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
int withdist = 0, withhash = 0, withcoords = 0;
int frommember = 0, fromloc = 0, byradius = 0, bybox = 0;
int sort = SORT_NONE;
long long count = 0;
int any = 0; /* any=1 means a limited search, stop as soon as enough results were found. */
long long count = 0; /* Max number of results to return. 0 means unlimited. */
if (c->argc > base_args) {
int remaining = c->argc - base_args;
for (int i = 0; i < remaining; i++) {
@ -561,6 +565,8 @@ void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
withhash = 1;
} else if (!strcasecmp(arg, "withcoord")) {
withcoords = 1;
} else if (!strcasecmp(arg, "any")) {
any = 1;
} else if (!strcasecmp(arg, "asc")) {
sort = SORT_ASC;
} else if (!strcasecmp(arg, "desc")) {
@ -662,16 +668,23 @@ void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
return;
}
/* COUNT without ordering does not make much sense, force ASC
* ordering if COUNT was specified but no sorting was requested. */
if (count != 0 && sort == SORT_NONE) sort = SORT_ASC;
if (any && !count) {
addReplyErrorFormat(c, "the ANY argument requires COUNT argument");
return;
}
/* COUNT without ordering does not make much sense (we need to
* sort in order to return the closest N entries),
* force ASC ordering if COUNT was specified but no sorting was
* requested. Note that this is not needed for ANY option. */
if (count != 0 && sort == SORT_NONE && !any) sort = SORT_ASC;
/* Get all neighbor geohash boxes for our radius search */
GeoHashRadius georadius = geohashCalculateAreasByShapeWGS84(&shape);
/* Search the zset for all matching points */
geoArray *ga = geoArrayCreate();
membersOfAllNeighbors(zobj, georadius, &shape, ga);
membersOfAllNeighbors(zobj, georadius, &shape, ga, any ? count : 0);
/* If no matching results, the user gets an empty reply. */
if (ga->used == 0 && storekey == NULL) {

View File

@ -178,6 +178,19 @@ start_server {tags {"geo"}} {
r georadius nyc -73.9798091 40.7598464 10 km COUNT 3
} {{central park n/q/r} 4545 {union square}}
test {GEORADIUS with ANY not sorted by default} {
r georadius nyc -73.9798091 40.7598464 10 km COUNT 3 ANY
} {{wtc one} {union square} {central park n/q/r}}
test {GEORADIUS with ANY sorted by ASC} {
r georadius nyc -73.9798091 40.7598464 10 km COUNT 3 ANY ASC
} {{central park n/q/r} {union square} {wtc one}}
test {GEORADIUS with ANY but no COUNT} {
catch {r georadius nyc -73.9798091 40.7598464 10 km ANY ASC} e
set e
} {ERR*ANY*requires*COUNT*}
test {GEORADIUS with COUNT but missing integer argument} {
catch {r georadius nyc -73.9798091 40.7598464 10 km COUNT} e
set e