summaryrefslogtreecommitdiff
path: root/src/geo.c
diff options
context:
space:
mode:
Diffstat (limited to 'src/geo.c')
-rw-r--r--src/geo.c95
1 files changed, 61 insertions, 34 deletions
diff --git a/src/geo.c b/src/geo.c
index 2fd1cb21b..9c668288b 100644
--- a/src/geo.c
+++ b/src/geo.c
@@ -100,8 +100,8 @@ int extractLongLatOrReply(client *c, robj **argv, double *xy) {
}
if (xy[0] < GEO_LONG_MIN || xy[0] > GEO_LONG_MAX ||
xy[1] < GEO_LAT_MIN || xy[1] > GEO_LAT_MAX) {
- addReplySds(c, sdscatprintf(sdsempty(),
- "-ERR invalid longitude,latitude pair %f,%f\r\n",xy[0],xy[1]));
+ addReplyErrorFormat(c,
+ "-ERR invalid longitude,latitude pair %f,%f\r\n",xy[0],xy[1]);
return C_ERR;
}
return C_OK;
@@ -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;
@@ -428,31 +431,45 @@ static int sort_gp_desc(const void *a, const void *b) {
* Commands
* ==================================================================== */
-/* GEOADD key long lat name [long2 lat2 name2 ... longN latN nameN] */
+/* GEOADD key [CH] [NX|XX] long lat name [long2 lat2 name2 ... longN latN nameN] */
void geoaddCommand(client *c) {
- /* Check arguments number for sanity. */
- if ((c->argc - 2) % 3 != 0) {
+ int xx = 0, nx = 0, longidx = 2;
+ int i;
+
+ /* Parse options. At the end 'longidx' is set to the argument position
+ * of the longitude of the first element. */
+ while (longidx < c->argc) {
+ char *opt = c->argv[longidx]->ptr;
+ if (!strcasecmp(opt,"nx")) nx = 1;
+ else if (!strcasecmp(opt,"xx")) xx = 1;
+ else if (!strcasecmp(opt,"ch")) {}
+ else break;
+ longidx++;
+ }
+
+ if ((c->argc - longidx) % 3 || (xx && nx)) {
/* Need an odd number of arguments if we got this far... */
- addReplyError(c, "syntax error. Try GEOADD key [x1] [y1] [name1] "
- "[x2] [y2] [name2] ... ");
+ addReplyErrorObject(c,shared.syntaxerr);
return;
}
- int elements = (c->argc - 2) / 3;
- int argc = 2+elements*2; /* ZADD key score ele ... */
+ /* Set up the vector for calling ZADD. */
+ int elements = (c->argc - longidx) / 3;
+ int argc = longidx+elements*2; /* ZADD key [CH] [NX|XX] score ele ... */
robj **argv = zcalloc(argc*sizeof(robj*));
argv[0] = createRawStringObject("zadd",4);
- argv[1] = c->argv[1]; /* key */
- incrRefCount(argv[1]);
+ for (i = 1; i < longidx; i++) {
+ argv[i] = c->argv[i];
+ incrRefCount(argv[i]);
+ }
/* Create the argument vector to call ZADD in order to add all
* the score,value pairs to the requested zset, where score is actually
* an encoded version of lat,long. */
- int i;
for (i = 0; i < elements; i++) {
double xy[2];
- if (extractLongLatOrReply(c, (c->argv+2)+(i*3),xy) == C_ERR) {
+ if (extractLongLatOrReply(c, (c->argv+longidx)+(i*3),xy) == C_ERR) {
for (i = 0; i < argc; i++)
if (argv[i]) decrRefCount(argv[i]);
zfree(argv);
@@ -464,9 +481,9 @@ void geoaddCommand(client *c) {
geohashEncodeWGS84(xy[0], xy[1], GEO_STEP_MAX, &hash);
GeoHashFix52Bits bits = geohashAlign52Bits(hash);
robj *score = createObject(OBJ_STRING, sdsfromlonglong(bits));
- robj *val = c->argv[2 + i * 3 + 2];
- argv[2+i*2] = score;
- argv[3+i*2] = val;
+ robj *val = c->argv[longidx + i * 3 + 2];
+ argv[longidx+i*2] = score;
+ argv[longidx+1+i*2] = val;
incrRefCount(val);
}
@@ -486,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] [FORMLOG long lat] [BYRADIUS radius unit]
- * [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT count] [ASC|DESC]
- * GEOSEARCHSTORE dest_key src_key [FROMMEMBER member] [FORMLOG long lat] [BYRADIUS radius unit]
- * [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT count] [ASC|DESC] [STOREDIST]
+ * GEOSEARCH key [FROMMEMBER member] [FROMLONLAT long lat] [BYRADIUS radius unit]
+ * [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 [ANY]] [ASC|DESC] [STOREDIST]
* */
void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
robj *storekey = NULL;
@@ -536,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++) {
@@ -547,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;
@@ -620,7 +640,7 @@ void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
bybox = 1;
i += 3;
} else {
- addReply(c, shared.syntaxerr);
+ addReplyErrorObject(c,shared.syntaxerr);
return;
}
}
@@ -648,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) {
@@ -902,7 +929,7 @@ void geodistCommand(client *c) {
to_meter = extractUnitOrReply(c,c->argv[4]);
if (to_meter < 0) return;
} else if (c->argc > 5) {
- addReply(c,shared.syntaxerr);
+ addReplyErrorObject(c,shared.syntaxerr);
return;
}