Add GEOSEARCH / GEOSEARCHSTORE commands (#8094)

Add commands to query geospatial data with bounding box.

Two new commands that replace the existing 4 GEORADIUS* commands.

GEOSEARCH key [FROMMEMBER member] [FROMLOC long lat] [BYRADIUS radius
unit] [BYBOX width height unit] [WITHCORD] [WITHDIST] [WITHASH] [COUNT
count] [ASC|DESC]

GEOSEARCHSTORE dest_key src_key [FROMMEMBER member] [FROMLOC long lat]
[BYRADIUS radius unit] [BYBOX width height unit] [WITHCORD] [WITHDIST]
[WITHASH] [COUNT count] [ASC|DESC] [STOREDIST]

- Add two types of CIRCULAR_TYPE and RECTANGLE_TYPE to achieve different searches
- Judge whether the point is within the rectangle, refer to:
geohashGetDistanceIfInRectangle

1 files changed, 155 insertions, 56 deletions

diff --git a/src/geo.c b/src/geo.c
index 7594bb2e9..13cbb1c50 100644
--- a/src/geo.c
+++ b/src/geo.c
@@ -144,31 +144,57 @@ double extractUnitOrReply(client *c, robj *unit) {
 
 /* Input Argument Helper.
  * Extract the distance from the specified two arguments starting at 'argv'
- * that should be in the form: <number> <unit>, and return the distance in the
- * specified unit on success. *conversions is populated with the coefficient
- * to use in order to convert meters to the unit.
- *
- * On error a value less than zero is returned. */
-double extractDistanceOrReply(client *c, robj **argv,
-                                     double *conversion) {
+ * that should be in the form: <number> <unit>, and return C_OK or C_ERR means success or failure
+ * *conversions is populated with the coefficient to use in order to convert meters to the unit.*/
+int extractDistanceOrReply(client *c, robj **argv,
+                              double *conversion, double *radius) {
     double distance;
     if (getDoubleFromObjectOrReply(c, argv[0], &distance,
                                    "need numeric radius") != C_OK) {
-        return -1;
+        return C_ERR;
     }
 
     if (distance < 0) {
         addReplyError(c,"radius cannot be negative");
-        return -1;
+        return C_ERR;
     }
+    if (radius) *radius = distance;
 
     double to_meters = extractUnitOrReply(c,argv[1]);
     if (to_meters < 0) {
-        return -1;
+        return C_ERR;
+    }
+
+    if (conversion) *conversion = to_meters;
+    return C_OK;
+}
+
+/* Input Argument Helper.
+ * Extract height and width from the specified three arguments starting at 'argv'
+ * that should be in the form: <number> <number> <unit>, and return C_OK or C_ERR means success or failure
+ * *conversions is populated with the coefficient to use in order to convert meters to the unit.*/
+int extractBoxOrReply(client *c, robj **argv, double *conversion,
+                         double *height, double *width) {
+    double h, w;
+    if ((getDoubleFromObjectOrReply(c, argv[0], &h, "need numeric height") != C_OK) ||
+        (getDoubleFromObjectOrReply(c, argv[1], &w, "need numeric width") != C_OK)) {
+        return C_ERR;
+    }
+
+    if (h < 0 || w < 0) {
+        addReplyError(c, "height or width cannot be negative");
+        return C_ERR;
+    }
+    if (height) *height = h;
+    if (width) *width = w;
+
+    double to_meters = extractUnitOrReply(c,argv[2]);
+    if (to_meters < 0) {
+        return C_ERR;
     }
 
     if (conversion) *conversion = to_meters;
-    return distance * to_meters;
+    return C_OK;
 }
 
 /* The default addReplyDouble has too much accuracy.  We use this
@@ -183,21 +209,22 @@ void addReplyDoubleDistance(client *c, double d) {
 }
 
 /* Helper function for geoGetPointsInRange(): given a sorted set score
- * representing a point, and another point (the center of our search) and
- * a radius, appends this entry as a geoPoint into the specified geoArray
- * only if the point is within the search area.
+ * representing a point, and a GeoShape, appends this entry as a geoPoint
+ * into the specified geoArray only if the point is within the search area.
  *
  * returns C_OK if the point is included, or REIDS_ERR if it is outside. */
-int geoAppendIfWithinRadius(geoArray *ga, double lon, double lat, double radius, double score, sds member) {
-    double distance, xy[2];
+int geoAppendIfWithinShape(geoArray *ga, GeoShape *shape, double score, sds member) {
+    double distance = 0, xy[2];
 
     if (!decodeGeohash(score,xy)) return C_ERR; /* Can't decode. */
     /* Note that geohashGetDistanceIfInRadiusWGS84() takes arguments in
      * reverse order: longitude first, latitude later. */
-    if (!geohashGetDistanceIfInRadiusWGS84(lon,lat, xy[0], xy[1],
-                                           radius, &distance))
-    {
-        return C_ERR;
+    if (shape->type == CIRCULAR_TYPE) {
+        if (!geohashGetDistanceIfInRadiusWGS84(shape->xy[0], shape->xy[1], xy[0], xy[1],
+                                               shape->t.radius*shape->conversion, &distance)) return C_ERR;
+    } else if (shape->type == RECTANGLE_TYPE) {
+        if (!geohashGetDistanceIfInRectangle(shape->bounds, shape->xy[0], shape->xy[1],
+                                             xy[0], xy[1], &distance)) return C_ERR;
     }
 
     /* Append the new element. */
@@ -222,7 +249,7 @@ int geoAppendIfWithinRadius(geoArray *ga, double lon, double lat, double radius,
  * 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, double lon, double lat, double radius, geoArray *ga) {
+int geoGetPointsInRange(robj *zobj, double min, double max, GeoShape *shape, geoArray *ga) {
     /* 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 };
@@ -254,7 +281,7 @@ int geoGetPointsInRange(robj *zobj, double min, double max, double lon, double l
             ziplistGet(eptr, &vstr, &vlen, &vlong);
             member = (vstr == NULL) ? sdsfromlonglong(vlong) :
                                       sdsnewlen(vstr,vlen);
-            if (geoAppendIfWithinRadius(ga,lon,lat,radius,score,member)
+            if (geoAppendIfWithinShape(ga,shape,score,member)
                 == C_ERR) sdsfree(member);
             zzlNext(zl, &eptr, &sptr);
         }
@@ -275,7 +302,7 @@ int geoGetPointsInRange(robj *zobj, double min, double max, double lon, double l
                 break;
 
             ele = sdsdup(ele);
-            if (geoAppendIfWithinRadius(ga,lon,lat,radius,ln->score,ele)
+            if (geoAppendIfWithinShape(ga,shape,ln->score,ele)
                 == C_ERR) sdsfree(ele);
             ln = ln->level[0].forward;
         }
@@ -315,15 +342,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, double lon, double lat, double radius) {
+int membersOfGeoHashBox(robj *zobj, GeoHashBits hash, geoArray *ga, GeoShape *shape) {
     GeoHashFix52Bits min, max;
 
     scoresOfGeoHashBox(hash,&min,&max);
-    return geoGetPointsInRange(zobj, min, max, lon, lat, radius, ga);
+    return geoGetPointsInRange(zobj, min, max, shape, ga);
 }
 
 /* Search all eight neighbors + self geohash box */
-int membersOfAllNeighbors(robj *zobj, GeoHashRadius n, double lon, double lat, double radius, geoArray *ga) {
+int membersOfAllNeighbors(robj *zobj, GeoHashRadius n, GeoShape *shape, geoArray *ga) {
     GeoHashBits neighbors[9];
     unsigned int i, count = 0, last_processed = 0;
     int debugmsg = 0;
@@ -374,7 +401,7 @@ int membersOfAllNeighbors(robj *zobj, GeoHashRadius n, double lon, double lat, d
                 D("Skipping processing of %d, same as previous\n",i);
             continue;
         }
-        count += membersOfGeoHashBox(zobj, neighbors[i], ga, lon, lat, radius);
+        count += membersOfGeoHashBox(zobj, neighbors[i], ga, shape);
         last_processed = i;
     }
     return count;
@@ -454,51 +481,60 @@ void geoaddCommand(client *c) {
 
 #define RADIUS_COORDS (1<<0)    /* Search around coordinates. */
 #define RADIUS_MEMBER (1<<1)    /* Search around member. */
-#define RADIUS_NOSTORE (1<<2)   /* Do not acceot STORE/STOREDIST option. */
+#define RADIUS_NOSTORE (1<<2)   /* Do not accept STORE/STOREDIST option. */
+#define GEOSEARCH (1<<3)        /* GEOSEARCH command variant (different arguments supported) */
+#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]
- * GEORADIUSBYMEMBER key member radius unit ... options ... */
-void georadiusGeneric(client *c, int flags) {
-    robj *key = c->argv[1];
+ * 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]
+ *  */
+void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
     robj *storekey = NULL;
     int storedist = 0; /* 0 for STORE, 1 for STOREDIST. */
 
     /* Look up the requested zset */
     robj *zobj = NULL;
-    if ((zobj = lookupKeyReadOrReply(c, key, shared.emptyarray)) == NULL ||
+    if ((zobj = lookupKeyReadOrReply(c, c->argv[srcKeyIndex], shared.emptyarray)) == NULL ||
         checkType(c, zobj, OBJ_ZSET)) {
         return;
     }
 
-    /* Find long/lat to use for radius search based on inquiry type */
+    /* Find long/lat to use for radius or box search based on inquiry type */
     int base_args;
-    double xy[2] = { 0 };
+    GeoShape shape = {0};
     if (flags & RADIUS_COORDS) {
         base_args = 6;
-        if (extractLongLatOrReply(c, c->argv + 2, xy) == C_ERR)
-            return;
+        shape.type = CIRCULAR_TYPE;
+        if (extractLongLatOrReply(c, c->argv + 2, shape.xy) == C_ERR) return;
+        if (extractDistanceOrReply(c, c->argv+base_args-2, &shape.conversion, &shape.t.radius) != C_OK) return;
     } else if (flags & RADIUS_MEMBER) {
         base_args = 5;
+        shape.type = CIRCULAR_TYPE;
         robj *member = c->argv[2];
-        if (longLatFromMember(zobj, member, xy) == C_ERR) {
+        if (longLatFromMember(zobj, member, shape.xy) == C_ERR) {
             addReplyError(c, "could not decode requested zset member");
             return;
         }
+        if (extractDistanceOrReply(c, c->argv+base_args-2, &shape.conversion, &shape.t.radius) != C_OK) return;
+    } else if (flags & GEOSEARCH) {
+        base_args = 2;
+        if (flags & GEOSEARCHSTORE) {
+            base_args = 3;
+            storekey = c->argv[1];
+        }
     } else {
         addReplyError(c, "Unknown georadius search type");
         return;
     }
 
-    /* Extract radius and units from arguments */
-    double radius_meters = 0, conversion = 1;
-    if ((radius_meters = extractDistanceOrReply(c, c->argv + base_args - 2,
-                                                &conversion)) < 0) {
-        return;
-    }
-
     /* Discover and populate all optional parameters. */
     int withdist = 0, withhash = 0, withcoords = 0;
+    int frommember = 0, fromloc = 0, byradius = 0, bybox = 0;
     int sort = SORT_NONE;
     long long count = 0;
     if (c->argc > base_args) {
@@ -525,18 +561,64 @@ void georadiusGeneric(client *c, int flags) {
                 i++;
             } else if (!strcasecmp(arg, "store") &&
                        (i+1) < remaining &&
-                       !(flags & RADIUS_NOSTORE))
+                       !(flags & RADIUS_NOSTORE) &&
+                       !(flags & GEOSEARCH))
             {
                 storekey = c->argv[base_args+i+1];
                 storedist = 0;
                 i++;
             } else if (!strcasecmp(arg, "storedist") &&
                        (i+1) < remaining &&
-                       !(flags & RADIUS_NOSTORE))
+                       !(flags & RADIUS_NOSTORE) &&
+                       !(flags & GEOSEARCH))
             {
                 storekey = c->argv[base_args+i+1];
                 storedist = 1;
                 i++;
+            } else if (!strcasecmp(arg, "storedist") &&
+                       (flags & GEOSEARCH) &&
+                       (flags & GEOSEARCHSTORE))
+            {
+                storedist = 1;
+            } else if (!strcasecmp(arg, "frommember") &&
+                      (i+1) < remaining &&
+                      flags & GEOSEARCH &&
+                      !fromloc)
+            {
+                if (longLatFromMember(zobj, c->argv[base_args+i+1], shape.xy) == C_ERR) {
+                    addReplyError(c, "could not decode requested zset member");
+                    return;
+                }
+                frommember = 1;
+                i++;
+            } else if (!strcasecmp(arg, "fromloc") &&
+                       (i+2) < remaining &&
+                       flags & GEOSEARCH &&
+                       !frommember)
+            {
+                if (extractLongLatOrReply(c, c->argv+base_args+i+1, shape.xy) == C_ERR) return;
+                fromloc = 1;
+                i += 2;
+            } else if (!strcasecmp(arg, "byradius") &&
+                       (i+2) < remaining &&
+                       flags & GEOSEARCH &&
+                       !bybox)
+            {
+                if (extractDistanceOrReply(c, c->argv+base_args+i+1, &shape.conversion, &shape.t.radius) != C_OK)
+                    return;
+                shape.type = CIRCULAR_TYPE;
+                byradius = 1;
+                i += 2;
+            } else if (!strcasecmp(arg, "bybox") &&
+                       (i+3) < remaining &&
+                       flags & GEOSEARCH &&
+                       !byradius)
+            {
+                if (extractBoxOrReply(c, c->argv+base_args+i+1, &shape.conversion, &shape.t.r.height,
+                        &shape.t.r.width) != C_OK) return;
+                shape.type = RECTANGLE_TYPE;
+                bybox = 1;
+                i += 3;
             } else {
                 addReply(c, shared.syntaxerr);
                 return;
@@ -552,17 +634,26 @@ void georadiusGeneric(client *c, int flags) {
         return;
     }
 
+    if ((flags & GEOSEARCH) && !(frommember || fromloc)) {
+        addReplyErrorFormat(c, "exactly one of FROMMEMBER or FROMLOC can be specified for %s", (char *)c->argv[0]->ptr);
+        return;
+    }
+
+    if ((flags & GEOSEARCH) && !(byradius || bybox)) {
+        addReplyErrorFormat(c, "exactly one of BYRADIUS and BYBOX can be specified for %s", (char *)c->argv[0]->ptr);
+        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;
 
     /* Get all neighbor geohash boxes for our radius search */
-    GeoHashRadius georadius =
-        geohashGetAreasByRadiusWGS84(xy[0], xy[1], radius_meters);
+    GeoHashRadius georadius = geohashCalculateAreasByShapeWGS84(&shape);
 
     /* Search the zset for all matching points */
     geoArray *ga = geoArrayCreate();
-    membersOfAllNeighbors(zobj, georadius, xy[0], xy[1], radius_meters, ga);
+    membersOfAllNeighbors(zobj, georadius, &shape, ga);
 
     /* If no matching results, the user gets an empty reply. */
     if (ga->used == 0 && storekey == NULL) {
@@ -607,7 +698,7 @@ void georadiusGeneric(client *c, int flags) {
         int i;
         for (i = 0; i < returned_items; i++) {
             geoPoint *gp = ga->array+i;
-            gp->dist /= conversion; /* Fix according to unit. */
+            gp->dist /= shape.conversion; /* Fix according to unit. */
 
             /* If we have options in option_length, return each sub-result
              * as a nested multi-bulk.  Add 1 to account for result value
@@ -645,7 +736,7 @@ void georadiusGeneric(client *c, int flags) {
         for (i = 0; i < returned_items; i++) {
             zskiplistNode *znode;
             geoPoint *gp = ga->array+i;
-            gp->dist /= conversion; /* Fix according to unit. */
+            gp->dist /= shape.conversion; /* Fix according to unit. */
             double score = storedist ? gp->dist : gp->score;
             size_t elelen = sdslen(gp->member);
 
@@ -659,7 +750,7 @@ void georadiusGeneric(client *c, int flags) {
             zsetConvertToZiplistIfNeeded(zobj,maxelelen);
             setKey(c,c->db,storekey,zobj);
             decrRefCount(zobj);
-            notifyKeyspaceEvent(NOTIFY_ZSET,"georadiusstore",storekey,
+            notifyKeyspaceEvent(NOTIFY_ZSET,flags & GEOSEARCH ? "geosearchstore" : "georadiusstore",storekey,
                                 c->db->id);
             server.dirty += returned_items;
         } else if (dbDelete(c->db,storekey)) {
@@ -674,22 +765,30 @@ void georadiusGeneric(client *c, int flags) {
 
 /* GEORADIUS wrapper function. */
 void georadiusCommand(client *c) {
-    georadiusGeneric(c, RADIUS_COORDS);
+    georadiusGeneric(c, 1, RADIUS_COORDS);
 }
 
 /* GEORADIUSBYMEMBER wrapper function. */
 void georadiusbymemberCommand(client *c) {
-    georadiusGeneric(c, RADIUS_MEMBER);
+    georadiusGeneric(c, 1, RADIUS_MEMBER);
 }
 
 /* GEORADIUS_RO wrapper function. */
 void georadiusroCommand(client *c) {
-    georadiusGeneric(c, RADIUS_COORDS|RADIUS_NOSTORE);
+    georadiusGeneric(c, 1, RADIUS_COORDS|RADIUS_NOSTORE);
 }
 
 /* GEORADIUSBYMEMBER_RO wrapper function. */
 void georadiusbymemberroCommand(client *c) {
-    georadiusGeneric(c, RADIUS_MEMBER|RADIUS_NOSTORE);
+    georadiusGeneric(c, 1, RADIUS_MEMBER|RADIUS_NOSTORE);
+}
+
+void geosearchCommand(client *c) {
+    georadiusGeneric(c, 1, GEOSEARCH);
+}
+
+void geosearchstoreCommand(client *c) {
+    georadiusGeneric(c, 2, GEOSEARCH|GEOSEARCHSTORE);
 }
 
 /* GEOHASH key ele1 ele2 ... eleN