diff --git a/src/geo.c b/src/geo.c index e9c8bf5cb..9c668288b 100644 --- a/src/geo.c +++ b/src/geo.c @@ -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,13 +565,15 @@ 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")) { sort = SORT_DESC; } else if (!strcasecmp(arg, "count") && (i+1) < remaining) { if (getLongLongFromObjectOrReply(c, c->argv[base_args+i+1], - &count, NULL) != C_OK) return; + &count, NULL) != C_OK) return; if (count <= 0) { addReplyError(c,"COUNT must be > 0"); return; @@ -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) { diff --git a/tests/unit/geo.tcl b/tests/unit/geo.tcl index 6c1f2ef11..0ac95f7c5 100644 --- a/tests/unit/geo.tcl +++ b/tests/unit/geo.tcl @@ -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