diff options
Diffstat (limited to 'src/geo.c')
-rw-r--r-- | src/geo.c | 95 |
1 files changed, 61 insertions, 34 deletions
@@ -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; } |