1 | /* |
2 | * Copyright (c) 2014, Matt Stancliff <[email protected]>. |
3 | * Copyright (c) 2015-2016, Salvatore Sanfilippo <[email protected]>. |
4 | * All rights reserved. |
5 | * |
6 | * Redistribution and use in source and binary forms, with or without |
7 | * modification, are permitted provided that the following conditions are met: |
8 | * |
9 | * * Redistributions of source code must retain the above copyright notice, |
10 | * this list of conditions and the following disclaimer. |
11 | * * Redistributions in binary form must reproduce the above copyright |
12 | * notice, this list of conditions and the following disclaimer in the |
13 | * documentation and/or other materials provided with the distribution. |
14 | * * Neither the name of Redis nor the names of its contributors may be used |
15 | * to endorse or promote products derived from this software without |
16 | * specific prior written permission. |
17 | * |
18 | * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" |
19 | * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE |
20 | * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE |
21 | * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE |
22 | * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR |
23 | * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF |
24 | * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS |
25 | * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN |
26 | * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) |
27 | * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE |
28 | * POSSIBILITY OF SUCH DAMAGE. |
29 | */ |
30 | |
31 | #include "geo.h" |
32 | #include "geohash_helper.h" |
33 | #include "debugmacro.h" |
34 | #include "pqsort.h" |
35 | |
36 | /* Things exported from t_zset.c only for geo.c, since it is the only other |
37 | * part of Redis that requires close zset introspection. */ |
38 | unsigned char *zzlFirstInRange(unsigned char *zl, zrangespec *range); |
39 | int zslValueLteMax(double value, zrangespec *spec); |
40 | |
41 | /* ==================================================================== |
42 | * This file implements the following commands: |
43 | * |
44 | * - geoadd - add coordinates for value to geoset |
45 | * - georadius - search radius by coordinates in geoset |
46 | * - georadiusbymember - search radius based on geoset member position |
47 | * ==================================================================== */ |
48 | |
49 | /* ==================================================================== |
50 | * geoArray implementation |
51 | * ==================================================================== */ |
52 | |
53 | /* Create a new array of geoPoints. */ |
54 | geoArray *geoArrayCreate(void) { |
55 | geoArray *ga = zmalloc(sizeof(*ga)); |
56 | /* It gets allocated on first geoArrayAppend() call. */ |
57 | ga->array = NULL; |
58 | ga->buckets = 0; |
59 | ga->used = 0; |
60 | return ga; |
61 | } |
62 | |
63 | /* Add a new entry and return its pointer so that the caller can populate |
64 | * it with data. */ |
65 | geoPoint *geoArrayAppend(geoArray *ga) { |
66 | if (ga->used == ga->buckets) { |
67 | ga->buckets = (ga->buckets == 0) ? 8 : ga->buckets*2; |
68 | ga->array = zrealloc(ga->array,sizeof(geoPoint)*ga->buckets); |
69 | } |
70 | geoPoint *gp = ga->array+ga->used; |
71 | ga->used++; |
72 | return gp; |
73 | } |
74 | |
75 | /* Destroy a geoArray created with geoArrayCreate(). */ |
76 | void geoArrayFree(geoArray *ga) { |
77 | size_t i; |
78 | for (i = 0; i < ga->used; i++) sdsfree(ga->array[i].member); |
79 | zfree(ga->array); |
80 | zfree(ga); |
81 | } |
82 | |
83 | /* ==================================================================== |
84 | * Helpers |
85 | * ==================================================================== */ |
86 | int decodeGeohash(double bits, double *xy) { |
87 | GeoHashBits hash = { .bits = (uint64_t)bits, .step = GEO_STEP_MAX }; |
88 | return geohashDecodeToLongLatWGS84(hash, xy); |
89 | } |
90 | |
91 | /* Input Argument Helper */ |
92 | /* Take a pointer to the latitude arg then use the next arg for longitude. |
93 | * On parse error C_ERR is returned, otherwise C_OK. */ |
94 | int (client *c, robj **argv, double *xy) { |
95 | int i; |
96 | for (i = 0; i < 2; i++) { |
97 | if (getDoubleFromObjectOrReply(c, argv[i], xy + i, NULL) != |
98 | C_OK) { |
99 | return C_ERR; |
100 | } |
101 | } |
102 | if (xy[0] < GEO_LONG_MIN || xy[0] > GEO_LONG_MAX || |
103 | xy[1] < GEO_LAT_MIN || xy[1] > GEO_LAT_MAX) { |
104 | addReplyErrorFormat(c, |
105 | "-ERR invalid longitude,latitude pair %f,%f\r\n" ,xy[0],xy[1]); |
106 | return C_ERR; |
107 | } |
108 | return C_OK; |
109 | } |
110 | |
111 | /* Input Argument Helper */ |
112 | /* Decode lat/long from a zset member's score. |
113 | * Returns C_OK on successful decoding, otherwise C_ERR is returned. */ |
114 | int longLatFromMember(robj *zobj, robj *member, double *xy) { |
115 | double score = 0; |
116 | |
117 | if (zsetScore(zobj, member->ptr, &score) == C_ERR) return C_ERR; |
118 | if (!decodeGeohash(score, xy)) return C_ERR; |
119 | return C_OK; |
120 | } |
121 | |
122 | /* Check that the unit argument matches one of the known units, and returns |
123 | * the conversion factor to meters (you need to divide meters by the conversion |
124 | * factor to convert to the right unit). |
125 | * |
126 | * If the unit is not valid, an error is reported to the client, and a value |
127 | * less than zero is returned. */ |
128 | double (client *c, robj *unit) { |
129 | char *u = unit->ptr; |
130 | |
131 | if (!strcasecmp(u, "m" )) { |
132 | return 1; |
133 | } else if (!strcasecmp(u, "km" )) { |
134 | return 1000; |
135 | } else if (!strcasecmp(u, "ft" )) { |
136 | return 0.3048; |
137 | } else if (!strcasecmp(u, "mi" )) { |
138 | return 1609.34; |
139 | } else { |
140 | addReplyError(c, |
141 | "unsupported unit provided. please use M, KM, FT, MI" ); |
142 | return -1; |
143 | } |
144 | } |
145 | |
146 | /* Input Argument Helper. |
147 | * Extract the distance from the specified two arguments starting at 'argv' |
148 | * that should be in the form: <number> <unit>, and return C_OK or C_ERR means success or failure |
149 | * *conversions is populated with the coefficient to use in order to convert meters to the unit.*/ |
150 | int (client *c, robj **argv, |
151 | double *conversion, double *radius) { |
152 | double distance; |
153 | if (getDoubleFromObjectOrReply(c, argv[0], &distance, |
154 | "need numeric radius" ) != C_OK) { |
155 | return C_ERR; |
156 | } |
157 | |
158 | if (distance < 0) { |
159 | addReplyError(c,"radius cannot be negative" ); |
160 | return C_ERR; |
161 | } |
162 | if (radius) *radius = distance; |
163 | |
164 | double to_meters = extractUnitOrReply(c,argv[1]); |
165 | if (to_meters < 0) { |
166 | return C_ERR; |
167 | } |
168 | |
169 | if (conversion) *conversion = to_meters; |
170 | return C_OK; |
171 | } |
172 | |
173 | /* Input Argument Helper. |
174 | * Extract height and width from the specified three arguments starting at 'argv' |
175 | * that should be in the form: <number> <number> <unit>, and return C_OK or C_ERR means success or failure |
176 | * *conversions is populated with the coefficient to use in order to convert meters to the unit.*/ |
177 | int (client *c, robj **argv, double *conversion, |
178 | double *width, double *height) { |
179 | double h, w; |
180 | if ((getDoubleFromObjectOrReply(c, argv[0], &w, "need numeric width" ) != C_OK) || |
181 | (getDoubleFromObjectOrReply(c, argv[1], &h, "need numeric height" ) != C_OK)) { |
182 | return C_ERR; |
183 | } |
184 | |
185 | if (h < 0 || w < 0) { |
186 | addReplyError(c, "height or width cannot be negative" ); |
187 | return C_ERR; |
188 | } |
189 | if (height) *height = h; |
190 | if (width) *width = w; |
191 | |
192 | double to_meters = extractUnitOrReply(c,argv[2]); |
193 | if (to_meters < 0) { |
194 | return C_ERR; |
195 | } |
196 | |
197 | if (conversion) *conversion = to_meters; |
198 | return C_OK; |
199 | } |
200 | |
201 | /* The default addReplyDouble has too much accuracy. We use this |
202 | * for returning location distances. "5.2145 meters away" is nicer |
203 | * than "5.2144992818115 meters away." We provide 4 digits after the dot |
204 | * so that the returned value is decently accurate even when the unit is |
205 | * the kilometer. */ |
206 | void addReplyDoubleDistance(client *c, double d) { |
207 | char dbuf[128]; |
208 | int dlen = snprintf(dbuf, sizeof(dbuf), "%.4f" , d); |
209 | addReplyBulkCBuffer(c, dbuf, dlen); |
210 | } |
211 | |
212 | /* Helper function for geoGetPointsInRange(): given a sorted set score |
213 | * representing a point, and a GeoShape, appends this entry as a geoPoint |
214 | * into the specified geoArray only if the point is within the search area. |
215 | * |
216 | * returns C_OK if the point is included, or C_ERR if it is outside. */ |
217 | int geoAppendIfWithinShape(geoArray *ga, GeoShape *shape, double score, sds member) { |
218 | double distance = 0, xy[2]; |
219 | |
220 | if (!decodeGeohash(score,xy)) return C_ERR; /* Can't decode. */ |
221 | /* Note that geohashGetDistanceIfInRadiusWGS84() takes arguments in |
222 | * reverse order: longitude first, latitude later. */ |
223 | if (shape->type == CIRCULAR_TYPE) { |
224 | if (!geohashGetDistanceIfInRadiusWGS84(shape->xy[0], shape->xy[1], xy[0], xy[1], |
225 | shape->t.radius*shape->conversion, &distance)) return C_ERR; |
226 | } else if (shape->type == RECTANGLE_TYPE) { |
227 | if (!geohashGetDistanceIfInRectangle(shape->t.r.width * shape->conversion, |
228 | shape->t.r.height * shape->conversion, |
229 | shape->xy[0], shape->xy[1], xy[0], xy[1], &distance)) |
230 | return C_ERR; |
231 | } |
232 | |
233 | /* Append the new element. */ |
234 | geoPoint *gp = geoArrayAppend(ga); |
235 | gp->longitude = xy[0]; |
236 | gp->latitude = xy[1]; |
237 | gp->dist = distance; |
238 | gp->member = member; |
239 | gp->score = score; |
240 | return C_OK; |
241 | } |
242 | |
243 | /* Query a Redis sorted set to extract all the elements between 'min' and |
244 | * 'max', appending them into the array of geoPoint structures 'geoArray'. |
245 | * The command returns the number of elements added to the array. |
246 | * |
247 | * Elements which are farther than 'radius' from the specified 'x' and 'y' |
248 | * coordinates are not included. |
249 | * |
250 | * The ability of this function to append to an existing set of points is |
251 | * important for good performances because querying by radius is performed |
252 | * using multiple queries to the sorted set, that we later need to sort |
253 | * via qsort. Similarly we need to be able to reject points outside the search |
254 | * radius area ASAP in order to allocate and process more points than needed. */ |
255 | int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geoArray *ga, unsigned long limit) { |
256 | /* minex 0 = include min in range; maxex 1 = exclude max in range */ |
257 | /* That's: min <= val < max */ |
258 | zrangespec range = { .min = min, .max = max, .minex = 0, .maxex = 1 }; |
259 | size_t origincount = ga->used; |
260 | sds member; |
261 | |
262 | if (zobj->encoding == OBJ_ENCODING_LISTPACK) { |
263 | unsigned char *zl = zobj->ptr; |
264 | unsigned char *eptr, *sptr; |
265 | unsigned char *vstr = NULL; |
266 | unsigned int vlen = 0; |
267 | long long vlong = 0; |
268 | double score = 0; |
269 | |
270 | if ((eptr = zzlFirstInRange(zl, &range)) == NULL) { |
271 | /* Nothing exists starting at our min. No results. */ |
272 | return 0; |
273 | } |
274 | |
275 | sptr = lpNext(zl, eptr); |
276 | while (eptr) { |
277 | score = zzlGetScore(sptr); |
278 | |
279 | /* If we fell out of range, break. */ |
280 | if (!zslValueLteMax(score, &range)) |
281 | break; |
282 | |
283 | vstr = lpGetValue(eptr, &vlen, &vlong); |
284 | member = (vstr == NULL) ? sdsfromlonglong(vlong) : |
285 | sdsnewlen(vstr,vlen); |
286 | if (geoAppendIfWithinShape(ga,shape,score,member) |
287 | == C_ERR) sdsfree(member); |
288 | if (ga->used && limit && ga->used >= limit) break; |
289 | zzlNext(zl, &eptr, &sptr); |
290 | } |
291 | } else if (zobj->encoding == OBJ_ENCODING_SKIPLIST) { |
292 | zset *zs = zobj->ptr; |
293 | zskiplist *zsl = zs->zsl; |
294 | zskiplistNode *ln; |
295 | |
296 | if ((ln = zslFirstInRange(zsl, &range)) == NULL) { |
297 | /* Nothing exists starting at our min. No results. */ |
298 | return 0; |
299 | } |
300 | |
301 | while (ln) { |
302 | sds ele = ln->ele; |
303 | /* Abort when the node is no longer in range. */ |
304 | if (!zslValueLteMax(ln->score, &range)) |
305 | break; |
306 | |
307 | ele = sdsdup(ele); |
308 | if (geoAppendIfWithinShape(ga,shape,ln->score,ele) |
309 | == C_ERR) sdsfree(ele); |
310 | if (ga->used && limit && ga->used >= limit) break; |
311 | ln = ln->level[0].forward; |
312 | } |
313 | } |
314 | return ga->used - origincount; |
315 | } |
316 | |
317 | /* Compute the sorted set scores min (inclusive), max (exclusive) we should |
318 | * query in order to retrieve all the elements inside the specified area |
319 | * 'hash'. The two scores are returned by reference in *min and *max. */ |
320 | void scoresOfGeoHashBox(GeoHashBits hash, GeoHashFix52Bits *min, GeoHashFix52Bits *max) { |
321 | /* We want to compute the sorted set scores that will include all the |
322 | * elements inside the specified Geohash 'hash', which has as many |
323 | * bits as specified by hash.step * 2. |
324 | * |
325 | * So if step is, for example, 3, and the hash value in binary |
326 | * is 101010, since our score is 52 bits we want every element which |
327 | * is in binary: 101010????????????????????????????????????????????? |
328 | * Where ? can be 0 or 1. |
329 | * |
330 | * To get the min score we just use the initial hash value left |
331 | * shifted enough to get the 52 bit value. Later we increment the |
332 | * 6 bit prefix (see the hash.bits++ statement), and get the new |
333 | * prefix: 101011, which we align again to 52 bits to get the maximum |
334 | * value (which is excluded from the search). So we get everything |
335 | * between the two following scores (represented in binary): |
336 | * |
337 | * 1010100000000000000000000000000000000000000000000000 (included) |
338 | * and |
339 | * 1010110000000000000000000000000000000000000000000000 (excluded). |
340 | */ |
341 | *min = geohashAlign52Bits(hash); |
342 | hash.bits++; |
343 | *max = geohashAlign52Bits(hash); |
344 | } |
345 | |
346 | /* Obtain all members between the min/max of this geohash bounding box. |
347 | * Populate a geoArray of GeoPoints by calling geoGetPointsInRange(). |
348 | * Return the number of points added to the array. */ |
349 | int membersOfGeoHashBox(robj *zobj, GeoHashBits hash, geoArray *ga, GeoShape *shape, unsigned long limit) { |
350 | GeoHashFix52Bits min, max; |
351 | |
352 | scoresOfGeoHashBox(hash,&min,&max); |
353 | return geoGetPointsInRange(zobj, min, max, shape, ga, limit); |
354 | } |
355 | |
356 | /* Search all eight neighbors + self geohash box */ |
357 | int membersOfAllNeighbors(robj *zobj, const GeoHashRadius *n, GeoShape *shape, geoArray *ga, unsigned long limit) { |
358 | GeoHashBits neighbors[9]; |
359 | unsigned int i, count = 0, last_processed = 0; |
360 | int debugmsg = 0; |
361 | |
362 | neighbors[0] = n->hash; |
363 | neighbors[1] = n->neighbors.north; |
364 | neighbors[2] = n->neighbors.south; |
365 | neighbors[3] = n->neighbors.east; |
366 | neighbors[4] = n->neighbors.west; |
367 | neighbors[5] = n->neighbors.north_east; |
368 | neighbors[6] = n->neighbors.north_west; |
369 | neighbors[7] = n->neighbors.south_east; |
370 | neighbors[8] = n->neighbors.south_west; |
371 | |
372 | /* For each neighbor (*and* our own hashbox), get all the matching |
373 | * members and add them to the potential result list. */ |
374 | for (i = 0; i < sizeof(neighbors) / sizeof(*neighbors); i++) { |
375 | if (HASHISZERO(neighbors[i])) { |
376 | if (debugmsg) D("neighbors[%d] is zero" ,i); |
377 | continue; |
378 | } |
379 | |
380 | /* Debugging info. */ |
381 | if (debugmsg) { |
382 | GeoHashRange long_range, lat_range; |
383 | geohashGetCoordRange(&long_range,&lat_range); |
384 | GeoHashArea myarea = {{0}}; |
385 | geohashDecode(long_range, lat_range, neighbors[i], &myarea); |
386 | |
387 | /* Dump center square. */ |
388 | D("neighbors[%d]:\n" ,i); |
389 | D("area.longitude.min: %f\n" , myarea.longitude.min); |
390 | D("area.longitude.max: %f\n" , myarea.longitude.max); |
391 | D("area.latitude.min: %f\n" , myarea.latitude.min); |
392 | D("area.latitude.max: %f\n" , myarea.latitude.max); |
393 | D("\n" ); |
394 | } |
395 | |
396 | /* When a huge Radius (in the 5000 km range or more) is used, |
397 | * adjacent neighbors can be the same, leading to duplicated |
398 | * elements. Skip every range which is the same as the one |
399 | * processed previously. */ |
400 | if (last_processed && |
401 | neighbors[i].bits == neighbors[last_processed].bits && |
402 | neighbors[i].step == neighbors[last_processed].step) |
403 | { |
404 | if (debugmsg) |
405 | D("Skipping processing of %d, same as previous\n" ,i); |
406 | continue; |
407 | } |
408 | if (ga->used && limit && ga->used >= limit) break; |
409 | count += membersOfGeoHashBox(zobj, neighbors[i], ga, shape, limit); |
410 | last_processed = i; |
411 | } |
412 | return count; |
413 | } |
414 | |
415 | /* Sort comparators for qsort() */ |
416 | static int sort_gp_asc(const void *a, const void *b) { |
417 | const struct geoPoint *gpa = a, *gpb = b; |
418 | /* We can't do adist - bdist because they are doubles and |
419 | * the comparator returns an int. */ |
420 | if (gpa->dist > gpb->dist) |
421 | return 1; |
422 | else if (gpa->dist == gpb->dist) |
423 | return 0; |
424 | else |
425 | return -1; |
426 | } |
427 | |
428 | static int sort_gp_desc(const void *a, const void *b) { |
429 | return -sort_gp_asc(a, b); |
430 | } |
431 | |
432 | /* ==================================================================== |
433 | * Commands |
434 | * ==================================================================== */ |
435 | |
436 | /* GEOADD key [CH] [NX|XX] long lat name [long2 lat2 name2 ... longN latN nameN] */ |
437 | void geoaddCommand(client *c) { |
438 | int xx = 0, nx = 0, longidx = 2; |
439 | int i; |
440 | |
441 | /* Parse options. At the end 'longidx' is set to the argument position |
442 | * of the longitude of the first element. */ |
443 | while (longidx < c->argc) { |
444 | char *opt = c->argv[longidx]->ptr; |
445 | if (!strcasecmp(opt,"nx" )) nx = 1; |
446 | else if (!strcasecmp(opt,"xx" )) xx = 1; |
447 | else if (!strcasecmp(opt,"ch" )) { /* Handle in zaddCommand. */ } |
448 | else break; |
449 | longidx++; |
450 | } |
451 | |
452 | if ((c->argc - longidx) % 3 || (xx && nx)) { |
453 | /* Need an odd number of arguments if we got this far... */ |
454 | addReplyErrorObject(c,shared.syntaxerr); |
455 | return; |
456 | } |
457 | |
458 | /* Set up the vector for calling ZADD. */ |
459 | int elements = (c->argc - longidx) / 3; |
460 | int argc = longidx+elements*2; /* ZADD key [CH] [NX|XX] score ele ... */ |
461 | robj **argv = zcalloc(argc*sizeof(robj*)); |
462 | argv[0] = createRawStringObject("zadd" ,4); |
463 | for (i = 1; i < longidx; i++) { |
464 | argv[i] = c->argv[i]; |
465 | incrRefCount(argv[i]); |
466 | } |
467 | |
468 | /* Create the argument vector to call ZADD in order to add all |
469 | * the score,value pairs to the requested zset, where score is actually |
470 | * an encoded version of lat,long. */ |
471 | for (i = 0; i < elements; i++) { |
472 | double xy[2]; |
473 | |
474 | if (extractLongLatOrReply(c, (c->argv+longidx)+(i*3),xy) == C_ERR) { |
475 | for (i = 0; i < argc; i++) |
476 | if (argv[i]) decrRefCount(argv[i]); |
477 | zfree(argv); |
478 | return; |
479 | } |
480 | |
481 | /* Turn the coordinates into the score of the element. */ |
482 | GeoHashBits hash; |
483 | geohashEncodeWGS84(xy[0], xy[1], GEO_STEP_MAX, &hash); |
484 | GeoHashFix52Bits bits = geohashAlign52Bits(hash); |
485 | robj *score = createObject(OBJ_STRING, sdsfromlonglong(bits)); |
486 | robj *val = c->argv[longidx + i * 3 + 2]; |
487 | argv[longidx+i*2] = score; |
488 | argv[longidx+1+i*2] = val; |
489 | incrRefCount(val); |
490 | } |
491 | |
492 | /* Finally call ZADD that will do the work for us. */ |
493 | replaceClientCommandVector(c,argc,argv); |
494 | zaddCommand(c); |
495 | } |
496 | |
497 | #define SORT_NONE 0 |
498 | #define SORT_ASC 1 |
499 | #define SORT_DESC 2 |
500 | |
501 | #define RADIUS_COORDS (1<<0) /* Search around coordinates. */ |
502 | #define RADIUS_MEMBER (1<<1) /* Search around member. */ |
503 | #define RADIUS_NOSTORE (1<<2) /* Do not accept STORE/STOREDIST option. */ |
504 | #define GEOSEARCH (1<<3) /* GEOSEARCH command variant (different arguments supported) */ |
505 | #define GEOSEARCHSTORE (1<<4) /* GEOSEARCHSTORE just accept STOREDIST option */ |
506 | |
507 | /* GEORADIUS key x y radius unit [WITHDIST] [WITHHASH] [WITHCOORD] [ASC|DESC] |
508 | * [COUNT count [ANY]] [STORE key] [STOREDIST key] |
509 | * GEORADIUSBYMEMBER key member radius unit ... options ... |
510 | * GEOSEARCH key [FROMMEMBER member] [FROMLONLAT long lat] [BYRADIUS radius unit] |
511 | * [BYBOX width height unit] [WITHCOORD] [WITHDIST] [WITHASH] [COUNT count [ANY]] [ASC|DESC] |
512 | * GEOSEARCHSTORE dest_key src_key [FROMMEMBER member] [FROMLONLAT long lat] [BYRADIUS radius unit] |
513 | * [BYBOX width height unit] [COUNT count [ANY]] [ASC|DESC] [STOREDIST] |
514 | * */ |
515 | void georadiusGeneric(client *c, int srcKeyIndex, int flags) { |
516 | robj *storekey = NULL; |
517 | int storedist = 0; /* 0 for STORE, 1 for STOREDIST. */ |
518 | |
519 | /* Look up the requested zset */ |
520 | robj *zobj = lookupKeyRead(c->db, c->argv[srcKeyIndex]); |
521 | if (checkType(c, zobj, OBJ_ZSET)) return; |
522 | |
523 | /* Find long/lat to use for radius or box search based on inquiry type */ |
524 | int base_args; |
525 | GeoShape shape = {0}; |
526 | if (flags & RADIUS_COORDS) { |
527 | /* GEORADIUS or GEORADIUS_RO */ |
528 | base_args = 6; |
529 | shape.type = CIRCULAR_TYPE; |
530 | if (extractLongLatOrReply(c, c->argv + 2, shape.xy) == C_ERR) return; |
531 | if (extractDistanceOrReply(c, c->argv+base_args-2, &shape.conversion, &shape.t.radius) != C_OK) return; |
532 | } else if ((flags & RADIUS_MEMBER) && !zobj) { |
533 | /* We don't have a source key, but we need to proceed with argument |
534 | * parsing, so we know which reply to use depending on the STORE flag. */ |
535 | base_args = 5; |
536 | } else if (flags & RADIUS_MEMBER) { |
537 | /* GEORADIUSBYMEMBER or GEORADIUSBYMEMBER_RO */ |
538 | base_args = 5; |
539 | shape.type = CIRCULAR_TYPE; |
540 | robj *member = c->argv[2]; |
541 | if (longLatFromMember(zobj, member, shape.xy) == C_ERR) { |
542 | addReplyError(c, "could not decode requested zset member" ); |
543 | return; |
544 | } |
545 | if (extractDistanceOrReply(c, c->argv+base_args-2, &shape.conversion, &shape.t.radius) != C_OK) return; |
546 | } else if (flags & GEOSEARCH) { |
547 | /* GEOSEARCH or GEOSEARCHSTORE */ |
548 | base_args = 2; |
549 | if (flags & GEOSEARCHSTORE) { |
550 | base_args = 3; |
551 | storekey = c->argv[1]; |
552 | } |
553 | } else { |
554 | addReplyError(c, "Unknown georadius search type" ); |
555 | return; |
556 | } |
557 | |
558 | /* Discover and populate all optional parameters. */ |
559 | int withdist = 0, withhash = 0, withcoords = 0; |
560 | int frommember = 0, fromloc = 0, byradius = 0, bybox = 0; |
561 | int sort = SORT_NONE; |
562 | int any = 0; /* any=1 means a limited search, stop as soon as enough results were found. */ |
563 | long long count = 0; /* Max number of results to return. 0 means unlimited. */ |
564 | if (c->argc > base_args) { |
565 | int remaining = c->argc - base_args; |
566 | for (int i = 0; i < remaining; i++) { |
567 | char *arg = c->argv[base_args + i]->ptr; |
568 | if (!strcasecmp(arg, "withdist" )) { |
569 | withdist = 1; |
570 | } else if (!strcasecmp(arg, "withhash" )) { |
571 | withhash = 1; |
572 | } else if (!strcasecmp(arg, "withcoord" )) { |
573 | withcoords = 1; |
574 | } else if (!strcasecmp(arg, "any" )) { |
575 | any = 1; |
576 | } else if (!strcasecmp(arg, "asc" )) { |
577 | sort = SORT_ASC; |
578 | } else if (!strcasecmp(arg, "desc" )) { |
579 | sort = SORT_DESC; |
580 | } else if (!strcasecmp(arg, "count" ) && (i+1) < remaining) { |
581 | if (getLongLongFromObjectOrReply(c, c->argv[base_args+i+1], |
582 | &count, NULL) != C_OK) return; |
583 | if (count <= 0) { |
584 | addReplyError(c,"COUNT must be > 0" ); |
585 | return; |
586 | } |
587 | i++; |
588 | } else if (!strcasecmp(arg, "store" ) && |
589 | (i+1) < remaining && |
590 | !(flags & RADIUS_NOSTORE) && |
591 | !(flags & GEOSEARCH)) |
592 | { |
593 | storekey = c->argv[base_args+i+1]; |
594 | storedist = 0; |
595 | i++; |
596 | } else if (!strcasecmp(arg, "storedist" ) && |
597 | (i+1) < remaining && |
598 | !(flags & RADIUS_NOSTORE) && |
599 | !(flags & GEOSEARCH)) |
600 | { |
601 | storekey = c->argv[base_args+i+1]; |
602 | storedist = 1; |
603 | i++; |
604 | } else if (!strcasecmp(arg, "storedist" ) && |
605 | (flags & GEOSEARCH) && |
606 | (flags & GEOSEARCHSTORE)) |
607 | { |
608 | storedist = 1; |
609 | } else if (!strcasecmp(arg, "frommember" ) && |
610 | (i+1) < remaining && |
611 | flags & GEOSEARCH && |
612 | !fromloc) |
613 | { |
614 | /* No source key, proceed with argument parsing and return an error when done. */ |
615 | if (zobj == NULL) { |
616 | frommember = 1; |
617 | i++; |
618 | continue; |
619 | } |
620 | |
621 | if (longLatFromMember(zobj, c->argv[base_args+i+1], shape.xy) == C_ERR) { |
622 | addReplyError(c, "could not decode requested zset member" ); |
623 | return; |
624 | } |
625 | frommember = 1; |
626 | i++; |
627 | } else if (!strcasecmp(arg, "fromlonlat" ) && |
628 | (i+2) < remaining && |
629 | flags & GEOSEARCH && |
630 | !frommember) |
631 | { |
632 | if (extractLongLatOrReply(c, c->argv+base_args+i+1, shape.xy) == C_ERR) return; |
633 | fromloc = 1; |
634 | i += 2; |
635 | } else if (!strcasecmp(arg, "byradius" ) && |
636 | (i+2) < remaining && |
637 | flags & GEOSEARCH && |
638 | !bybox) |
639 | { |
640 | if (extractDistanceOrReply(c, c->argv+base_args+i+1, &shape.conversion, &shape.t.radius) != C_OK) |
641 | return; |
642 | shape.type = CIRCULAR_TYPE; |
643 | byradius = 1; |
644 | i += 2; |
645 | } else if (!strcasecmp(arg, "bybox" ) && |
646 | (i+3) < remaining && |
647 | flags & GEOSEARCH && |
648 | !byradius) |
649 | { |
650 | if (extractBoxOrReply(c, c->argv+base_args+i+1, &shape.conversion, &shape.t.r.width, |
651 | &shape.t.r.height) != C_OK) return; |
652 | shape.type = RECTANGLE_TYPE; |
653 | bybox = 1; |
654 | i += 3; |
655 | } else { |
656 | addReplyErrorObject(c,shared.syntaxerr); |
657 | return; |
658 | } |
659 | } |
660 | } |
661 | |
662 | /* Trap options not compatible with STORE and STOREDIST. */ |
663 | if (storekey && (withdist || withhash || withcoords)) { |
664 | addReplyErrorFormat(c, |
665 | "%s is not compatible with WITHDIST, WITHHASH and WITHCOORD options" , |
666 | flags & GEOSEARCHSTORE? "GEOSEARCHSTORE" : "STORE option in GEORADIUS" ); |
667 | return; |
668 | } |
669 | |
670 | if ((flags & GEOSEARCH) && !(frommember || fromloc)) { |
671 | addReplyErrorFormat(c, |
672 | "exactly one of FROMMEMBER or FROMLONLAT can be specified for %s" , |
673 | (char *)c->argv[0]->ptr); |
674 | return; |
675 | } |
676 | |
677 | if ((flags & GEOSEARCH) && !(byradius || bybox)) { |
678 | addReplyErrorFormat(c, |
679 | "exactly one of BYRADIUS and BYBOX can be specified for %s" , |
680 | (char *)c->argv[0]->ptr); |
681 | return; |
682 | } |
683 | |
684 | if (any && !count) { |
685 | addReplyErrorFormat(c, "the ANY argument requires COUNT argument" ); |
686 | return; |
687 | } |
688 | |
689 | /* Return ASAP when src key does not exist. */ |
690 | if (zobj == NULL) { |
691 | if (storekey) { |
692 | /* store key is not NULL, try to delete it and return 0. */ |
693 | if (dbDelete(c->db, storekey)) { |
694 | signalModifiedKey(c, c->db, storekey); |
695 | notifyKeyspaceEvent(NOTIFY_GENERIC, "del" , storekey, c->db->id); |
696 | server.dirty++; |
697 | } |
698 | addReply(c, shared.czero); |
699 | } else { |
700 | /* Otherwise we return an empty array. */ |
701 | addReply(c, shared.emptyarray); |
702 | } |
703 | return; |
704 | } |
705 | |
706 | /* COUNT without ordering does not make much sense (we need to |
707 | * sort in order to return the closest N entries), |
708 | * force ASC ordering if COUNT was specified but no sorting was |
709 | * requested. Note that this is not needed for ANY option. */ |
710 | if (count != 0 && sort == SORT_NONE && !any) sort = SORT_ASC; |
711 | |
712 | /* Get all neighbor geohash boxes for our radius search */ |
713 | GeoHashRadius georadius = geohashCalculateAreasByShapeWGS84(&shape); |
714 | |
715 | /* Search the zset for all matching points */ |
716 | geoArray *ga = geoArrayCreate(); |
717 | membersOfAllNeighbors(zobj, &georadius, &shape, ga, any ? count : 0); |
718 | |
719 | /* If no matching results, the user gets an empty reply. */ |
720 | if (ga->used == 0 && storekey == NULL) { |
721 | addReply(c,shared.emptyarray); |
722 | geoArrayFree(ga); |
723 | return; |
724 | } |
725 | |
726 | long result_length = ga->used; |
727 | long returned_items = (count == 0 || result_length < count) ? |
728 | result_length : count; |
729 | long option_length = 0; |
730 | |
731 | /* Process [optional] requested sorting */ |
732 | if (sort != SORT_NONE) { |
733 | int (*sort_gp_callback)(const void *a, const void *b) = NULL; |
734 | if (sort == SORT_ASC) { |
735 | sort_gp_callback = sort_gp_asc; |
736 | } else if (sort == SORT_DESC) { |
737 | sort_gp_callback = sort_gp_desc; |
738 | } |
739 | |
740 | if (returned_items == result_length) { |
741 | qsort(ga->array, result_length, sizeof(geoPoint), sort_gp_callback); |
742 | } else { |
743 | pqsort(ga->array, result_length, sizeof(geoPoint), sort_gp_callback, |
744 | 0, (returned_items - 1)); |
745 | } |
746 | } |
747 | |
748 | if (storekey == NULL) { |
749 | /* No target key, return results to user. */ |
750 | |
751 | /* Our options are self-contained nested multibulk replies, so we |
752 | * only need to track how many of those nested replies we return. */ |
753 | if (withdist) |
754 | option_length++; |
755 | |
756 | if (withcoords) |
757 | option_length++; |
758 | |
759 | if (withhash) |
760 | option_length++; |
761 | |
762 | /* The array len we send is exactly result_length. The result is |
763 | * either all strings of just zset members *or* a nested multi-bulk |
764 | * reply containing the zset member string _and_ all the additional |
765 | * options the user enabled for this request. */ |
766 | addReplyArrayLen(c, returned_items); |
767 | |
768 | /* Finally send results back to the caller */ |
769 | int i; |
770 | for (i = 0; i < returned_items; i++) { |
771 | geoPoint *gp = ga->array+i; |
772 | gp->dist /= shape.conversion; /* Fix according to unit. */ |
773 | |
774 | /* If we have options in option_length, return each sub-result |
775 | * as a nested multi-bulk. Add 1 to account for result value |
776 | * itself. */ |
777 | if (option_length) |
778 | addReplyArrayLen(c, option_length + 1); |
779 | |
780 | addReplyBulkSds(c,gp->member); |
781 | gp->member = NULL; |
782 | |
783 | if (withdist) |
784 | addReplyDoubleDistance(c, gp->dist); |
785 | |
786 | if (withhash) |
787 | addReplyLongLong(c, gp->score); |
788 | |
789 | if (withcoords) { |
790 | addReplyArrayLen(c, 2); |
791 | addReplyHumanLongDouble(c, gp->longitude); |
792 | addReplyHumanLongDouble(c, gp->latitude); |
793 | } |
794 | } |
795 | } else { |
796 | /* Target key, create a sorted set with the results. */ |
797 | robj *zobj; |
798 | zset *zs; |
799 | int i; |
800 | size_t maxelelen = 0, totelelen = 0; |
801 | |
802 | if (returned_items) { |
803 | zobj = createZsetObject(); |
804 | zs = zobj->ptr; |
805 | } |
806 | |
807 | for (i = 0; i < returned_items; i++) { |
808 | zskiplistNode *znode; |
809 | geoPoint *gp = ga->array+i; |
810 | gp->dist /= shape.conversion; /* Fix according to unit. */ |
811 | double score = storedist ? gp->dist : gp->score; |
812 | size_t elelen = sdslen(gp->member); |
813 | |
814 | if (maxelelen < elelen) maxelelen = elelen; |
815 | totelelen += elelen; |
816 | znode = zslInsert(zs->zsl,score,gp->member); |
817 | serverAssert(dictAdd(zs->dict,gp->member,&znode->score) == DICT_OK); |
818 | gp->member = NULL; |
819 | } |
820 | |
821 | if (returned_items) { |
822 | zsetConvertToListpackIfNeeded(zobj,maxelelen,totelelen); |
823 | setKey(c,c->db,storekey,zobj,0); |
824 | decrRefCount(zobj); |
825 | notifyKeyspaceEvent(NOTIFY_ZSET,flags & GEOSEARCH ? "geosearchstore" : "georadiusstore" ,storekey, |
826 | c->db->id); |
827 | server.dirty += returned_items; |
828 | } else if (dbDelete(c->db,storekey)) { |
829 | signalModifiedKey(c,c->db,storekey); |
830 | notifyKeyspaceEvent(NOTIFY_GENERIC,"del" ,storekey,c->db->id); |
831 | server.dirty++; |
832 | } |
833 | addReplyLongLong(c, returned_items); |
834 | } |
835 | geoArrayFree(ga); |
836 | } |
837 | |
838 | /* GEORADIUS wrapper function. */ |
839 | void georadiusCommand(client *c) { |
840 | georadiusGeneric(c, 1, RADIUS_COORDS); |
841 | } |
842 | |
843 | /* GEORADIUSBYMEMBER wrapper function. */ |
844 | void georadiusbymemberCommand(client *c) { |
845 | georadiusGeneric(c, 1, RADIUS_MEMBER); |
846 | } |
847 | |
848 | /* GEORADIUS_RO wrapper function. */ |
849 | void georadiusroCommand(client *c) { |
850 | georadiusGeneric(c, 1, RADIUS_COORDS|RADIUS_NOSTORE); |
851 | } |
852 | |
853 | /* GEORADIUSBYMEMBER_RO wrapper function. */ |
854 | void georadiusbymemberroCommand(client *c) { |
855 | georadiusGeneric(c, 1, RADIUS_MEMBER|RADIUS_NOSTORE); |
856 | } |
857 | |
858 | void geosearchCommand(client *c) { |
859 | georadiusGeneric(c, 1, GEOSEARCH); |
860 | } |
861 | |
862 | void geosearchstoreCommand(client *c) { |
863 | georadiusGeneric(c, 2, GEOSEARCH|GEOSEARCHSTORE); |
864 | } |
865 | |
866 | /* GEOHASH key ele1 ele2 ... eleN |
867 | * |
868 | * Returns an array with an 11 characters geohash representation of the |
869 | * position of the specified elements. */ |
870 | void geohashCommand(client *c) { |
871 | char *geoalphabet= "0123456789bcdefghjkmnpqrstuvwxyz" ; |
872 | int j; |
873 | |
874 | /* Look up the requested zset */ |
875 | robj *zobj = lookupKeyRead(c->db, c->argv[1]); |
876 | if (checkType(c, zobj, OBJ_ZSET)) return; |
877 | |
878 | /* Geohash elements one after the other, using a null bulk reply for |
879 | * missing elements. */ |
880 | addReplyArrayLen(c,c->argc-2); |
881 | for (j = 2; j < c->argc; j++) { |
882 | double score; |
883 | if (!zobj || zsetScore(zobj, c->argv[j]->ptr, &score) == C_ERR) { |
884 | addReplyNull(c); |
885 | } else { |
886 | /* The internal format we use for geocoding is a bit different |
887 | * than the standard, since we use as initial latitude range |
888 | * -85,85, while the normal geohashing algorithm uses -90,90. |
889 | * So we have to decode our position and re-encode using the |
890 | * standard ranges in order to output a valid geohash string. */ |
891 | |
892 | /* Decode... */ |
893 | double xy[2]; |
894 | if (!decodeGeohash(score,xy)) { |
895 | addReplyNull(c); |
896 | continue; |
897 | } |
898 | |
899 | /* Re-encode */ |
900 | GeoHashRange r[2]; |
901 | GeoHashBits hash; |
902 | r[0].min = -180; |
903 | r[0].max = 180; |
904 | r[1].min = -90; |
905 | r[1].max = 90; |
906 | geohashEncode(&r[0],&r[1],xy[0],xy[1],26,&hash); |
907 | |
908 | char buf[12]; |
909 | int i; |
910 | for (i = 0; i < 11; i++) { |
911 | int idx; |
912 | if (i == 10) { |
913 | /* We have just 52 bits, but the API used to output |
914 | * an 11 bytes geohash. For compatibility we assume |
915 | * zero. */ |
916 | idx = 0; |
917 | } else { |
918 | idx = (hash.bits >> (52-((i+1)*5))) & 0x1f; |
919 | } |
920 | buf[i] = geoalphabet[idx]; |
921 | } |
922 | buf[11] = '\0'; |
923 | addReplyBulkCBuffer(c,buf,11); |
924 | } |
925 | } |
926 | } |
927 | |
928 | /* GEOPOS key ele1 ele2 ... eleN |
929 | * |
930 | * Returns an array of two-items arrays representing the x,y position of each |
931 | * element specified in the arguments. For missing elements NULL is returned. */ |
932 | void geoposCommand(client *c) { |
933 | int j; |
934 | |
935 | /* Look up the requested zset */ |
936 | robj *zobj = lookupKeyRead(c->db, c->argv[1]); |
937 | if (checkType(c, zobj, OBJ_ZSET)) return; |
938 | |
939 | /* Report elements one after the other, using a null bulk reply for |
940 | * missing elements. */ |
941 | addReplyArrayLen(c,c->argc-2); |
942 | for (j = 2; j < c->argc; j++) { |
943 | double score; |
944 | if (!zobj || zsetScore(zobj, c->argv[j]->ptr, &score) == C_ERR) { |
945 | addReplyNullArray(c); |
946 | } else { |
947 | /* Decode... */ |
948 | double xy[2]; |
949 | if (!decodeGeohash(score,xy)) { |
950 | addReplyNullArray(c); |
951 | continue; |
952 | } |
953 | addReplyArrayLen(c,2); |
954 | addReplyHumanLongDouble(c,xy[0]); |
955 | addReplyHumanLongDouble(c,xy[1]); |
956 | } |
957 | } |
958 | } |
959 | |
960 | /* GEODIST key ele1 ele2 [unit] |
961 | * |
962 | * Return the distance, in meters by default, otherwise according to "unit", |
963 | * between points ele1 and ele2. If one or more elements are missing NULL |
964 | * is returned. */ |
965 | void geodistCommand(client *c) { |
966 | double to_meter = 1; |
967 | |
968 | /* Check if there is the unit to extract, otherwise assume meters. */ |
969 | if (c->argc == 5) { |
970 | to_meter = extractUnitOrReply(c,c->argv[4]); |
971 | if (to_meter < 0) return; |
972 | } else if (c->argc > 5) { |
973 | addReplyErrorObject(c,shared.syntaxerr); |
974 | return; |
975 | } |
976 | |
977 | /* Look up the requested zset */ |
978 | robj *zobj = NULL; |
979 | if ((zobj = lookupKeyReadOrReply(c, c->argv[1], shared.null[c->resp])) |
980 | == NULL || checkType(c, zobj, OBJ_ZSET)) return; |
981 | |
982 | /* Get the scores. We need both otherwise NULL is returned. */ |
983 | double score1, score2, xyxy[4]; |
984 | if (zsetScore(zobj, c->argv[2]->ptr, &score1) == C_ERR || |
985 | zsetScore(zobj, c->argv[3]->ptr, &score2) == C_ERR) |
986 | { |
987 | addReplyNull(c); |
988 | return; |
989 | } |
990 | |
991 | /* Decode & compute the distance. */ |
992 | if (!decodeGeohash(score1,xyxy) || !decodeGeohash(score2,xyxy+2)) |
993 | addReplyNull(c); |
994 | else |
995 | addReplyDoubleDistance(c, |
996 | geohashGetDistance(xyxy[0],xyxy[1],xyxy[2],xyxy[3]) / to_meter); |
997 | } |
998 | |