Redis源码阅读之geo

Redis的geospatial（地理空间）功能是通过对geohash的封装来实现的。具体的底层实现源码位于redis/src/geo.c、redis/src/geohash_helper.c和redis/src/geohash.c文件中。

首先，让我们来看一下geoadd命令的实现，这个命令用于将地理位置添加到Redis中的有序集合中。

c
void geoaddCommand(client *c) {
    // 解析命令参数
    // ...

    // 计算地理位置的geohash值
    GeoHashBits hash;
    geohashEncodeWGS84(xy, xy[[1]](https://blog.csdn.net/weixin_46307478/article/details/123086299), GEO_STEP_MAX, &hash);
    GeoHashFix52Bits bits = geohashAlign52Bits(hash);
    robj *score = createObject(OBJ_STRING, sdsfromlonglong(bits));
    robj *val = c->argv[longidx + i * 3 + 2];
    argv[longidx+i*2] = score;
    argv[longidx+1+i*2] = val;
    incrRefCount(val);

    // 调用zadd命令将地理位置添加到有序集合中
    replaceClientCommandVector(c,argc,argv);
    zaddCommand(c);
}

可以看到，geoadd命令实际上是将地理位置的经纬度转换为geohash值，并将其作为分数（score）存储在有序集合中。

接下来，让我们来看一下georadius命令的实现，这个命令用于根据给定的地理位置和半径范围来搜索附近的地理位置。

c
void georadiusGeneric(client *c, int srcKeyIndex, int flags) {
    // 解析命令参数
    // ...

    // 获取有序集合对象
    robj *zobj = NULL;
    if ((zobj = lookupKeyReadOrReply(c, c->argv[srcKeyIndex], shared.emptyarray)) == NULL ||
        checkType(c, zobj, OBJ_ZSET)) {
        return;
    }

    // 创建一个geoArray对象
    geoArray *ga = geoArrayCreate();

    // 调用membersOfAllNeighbors函数获取附近的地理位置
    membersOfAllNeighbors(zobj, georadius, &shape, ga, any ? count : 0);

    // 处理搜索结果
    // ...

    geoArrayFree(ga);
}

在georadiusGeneric函数中，首先获取了存储地理位置的有序集合对象，然后调用membersOfAllNeighbors函数来获取附近的地理位置。最后，对搜索结果进行处理。

Redis的geospatial功能是通过对geohash的封装来实现的。实现细节可以参考源码中的geo.c、geohash_helper.c和geohash.c文件。

geo的结构

c
#ifndef __GEO_H__
#define __GEO_H__

#include "server.h"

/* 在 geo.c 中用于表示地球上的点的结构 */
typedef struct geoPoint {
    double longitude;   // 经度
    double latitude;    // 纬度
    double dist;        // 距离（未使用，可能是保留字段）
    double score;       // 分数（未使用，可能是保留字段）
    char *member;       // 成员名称（字符串）
} geoPoint;

/* 在 geo.c 中用于表示地球上的点数组的结构 */
typedef struct geoArray {
    struct geoPoint *array;  // 指向 geoPoint 结构数组的指针
    size_t buckets;          // 数组的桶数目（未使用，可能是保留字段）
    size_t used;             // 数组中已使用的元素数量
} geoArray;

#endif

CRUD函数

c
/* Create a new array of geoPoints. */
geoArray *geoArrayCreate(void) {
    geoArray *ga = zmalloc(sizeof(*ga));
    /* It gets allocated on first geoArrayAppend() call. */
    ga->array = NULL;
    ga->buckets = 0;
    ga->used = 0;
    return ga;
}

/* Add and populate with data a new entry to the geoArray. */
geoPoint *geoArrayAppend(geoArray *ga, double *xy, double dist,
                         double score, char *member)
{
    if (ga->used == ga->buckets) {
        ga->buckets = (ga->buckets == 0) ? 8 : ga->buckets*2;
        ga->array = zrealloc(ga->array,sizeof(geoPoint)*ga->buckets);
    }
    geoPoint *gp = ga->array+ga->used;
    gp->longitude = xy[0];
    gp->latitude = xy[1];
    gp->dist = dist;
    gp->member = member;
    gp->score = score;
    ga->used++;
    return gp;
}

/* Destroy a geoArray created with geoArrayCreate(). */
void geoArrayFree(geoArray *ga) {
    size_t i;
    for (i = 0; i < ga->used; i++) sdsfree(ga->array[i].member);
    zfree(ga->array);
    zfree(ga);
}

geoadd命令

c
/* GEOADD key [CH] [NX|XX] long lat name [long2 lat2 name2 ... longN latN nameN] */
void geoaddCommand(client *c) {
    int xx = 0, nx = 0, longidx = 2;
    int i;

    /* 解析选项。最后 'longidx' 被设置为第一个元素的经度的参数位置。 */
    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")) { /* 在 zaddCommand 中处理。 */ }
        else break;
        longidx++;
    }

    if ((c->argc - longidx) % 3 || (xx && nx)) {
        /* 如果执行到这里，需要奇数个参数... */
        addReplyErrorObject(c, shared.syntaxerr);
        return;
    }

    /* 为调用 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);
    for (i = 1; i < longidx; i++) {
        argv[i] = c->argv[i];
        incrRefCount(argv[i]);
    }

    /* 创建调用 ZADD 的参数向量，以将所有 score,value 对添加到请求的 zset 中，
     * 其中 score 实际上是 lat,long 的编码版本。 */
    for (i = 0; i < elements; i++) {
        double xy[2];

        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);
            return;
        }

        /* 将坐标转换为元素的分数。*/
        GeoHashBits hash;
        geohashEncodeWGS84(xy[0], xy[1], GEO_STEP_MAX, &hash);
        GeoHashFix52Bits bits = geohashAlign52Bits(hash);
        robj *score = createStringObjectFromLongLongWithSds(bits);
        robj *val = c->argv[longidx + i * 3 + 2];
        argv[longidx + i * 2] = score;
        argv[longidx + 1 + i * 2] = val;
        incrRefCount(val);
    }

    /* 最后调用 ZADD 完成工作。 */
    replaceClientCommandVector(c, argc, argv);
    zaddCommand(c);
}

其他命令

c
/* GEOPOS key ele1 ele2 ... eleN
 *
 * 返回一个包含每个指定元素的 x、y 位置的二维数组。对于缺失的元素，返回 NULL。*/
void geoposCommand(client *c) {
    int j;

    /* 查找请求的有序集合 */
    robj *zobj = lookupKeyRead(c->db, c->argv[1]);
    if (checkType(c, zobj, OBJ_ZSET)) return;

    /* 逐个报告元素的位置，对于缺失的元素使用 null 批量回复。*/
    addReplyArrayLen(c, c->argc - 2);
    for (j = 2; j < c->argc; j++) {
        double score;
        if (!zobj || zsetScore(zobj, c->argv[j]->ptr, &score) == C_ERR) {
            addReplyNullArray(c);
        } else {
            /* 解码... */
            double xy[2];
            if (!decodeGeohash(score, xy)) {
                addReplyNullArray(c);
                continue;
            }
            addReplyArrayLen(c, 2);
            addReplyHumanLongDouble(c, xy[0]);
            addReplyHumanLongDouble(c, xy[1]);
        }
    }
}

/* GEODIST key ele1 ele2 [unit]
 *
 * 返回 ele1 和 ele2 之间的距离，单位默认为米，否则根据 "unit" 参数确定。如果有一个或多个元素缺失，则返回 NULL。*/
void geodistCommand(client *c) {
    double to_meter = 1;

    /* 检查是否提供了单位，否则默认为米。 */
    if (c->argc == 5) {
        to_meter = extractUnitOrReply(c, c->argv[4]);
        if (to_meter < 0) return;
    } else if (c->argc > 5) {
        addReplyErrorObject(c, shared.syntaxerr);
        return;
    }

    /* 查找请求的有序集合 */
    robj *zobj = NULL;
    if ((zobj = lookupKeyReadOrReply(c, c->argv[1], shared.null[c->resp]))
        == NULL || checkType(c, zobj, OBJ_ZSET)) return;

    /* 获取分数。如果一个或两个元素缺失，则返回 NULL。 */
    double score1, score2, xyxy[4];
    if (zsetScore(zobj, c->argv[2]->ptr, &score1) == C_ERR ||
        zsetScore(zobj, c->argv[3]->ptr, &score2) == C_ERR)
    {
        addReplyNull(c);
        return;
    }

    /* 解码并计算距离。*/
    if (!decodeGeohash(score1, xyxy) || !decodeGeohash(score2, xyxy + 2))
        addReplyNull(c);
    else
        addReplyDoubleDistance(c,
            geohashGetDistance(xyxy[0], xyxy[1], xyxy[2], xyxy[3]) / to_meter);
}

geohash

c
/* GEOHASH key ele1 ele2 ... eleN
 *
 * 返回一个包含指定元素位置的 11 个字符 Geohash 表示的数组。*/
void geohashCommand(client *c) {
    char *geoalphabet = "0123456789bcdefghjkmnpqrstuvwxyz";
    int j;

    /* 查找请求的有序集合 */
    robj *zobj = lookupKeyRead(c->db, c->argv[1]);
    if (checkType(c, zobj, OBJ_ZSET)) return;

    /* 逐个 Geohash 元素，对于缺失的元素使用 null 批量回复。*/
    addReplyArrayLen(c, c->argc - 2);
    for (j = 2; j < c->argc; j++) {
        double score;
        if (!zobj || zsetScore(zobj, c->argv[j]->ptr, &score) == C_ERR) {
            addReplyNull(c);
        } else {
            /* 我们用于地理编码的内部格式与标准格式有点不同，
             * 因为我们的初始纬度范围是 -85,85，而标准的 Geohashing 算法使用 -90,90。
             * 因此，我们必须解码我们的位置并使用标准范围重新编码，以便输出有效的 Geohash 字符串。*/

            /* 解码... */
            double xy[2];
            if (!decodeGeohash(score, xy)) {
                addReplyNull(c);
                continue;
            }

            /* 重新编码 */
            GeoHashRange r[2];
            GeoHashBits hash;
            r[0].min = -180;
            r[0].max = 180;
            r[1].min = -90;
            r[1].max = 90;
            geohashEncode(&r[0], &r[1], xy[0], xy[1], 26, &hash);

            char buf[12];
            int i;
            for (i = 0; i < 11; i++) {
                int idx;
                if (i == 10) {
                    /* 我们只有 52 位，但是 API 用于输出一个 11 字节的 Geohash。
                     * 为了兼容性，我们假设为零。*/
                    idx = 0;
                } else {
                    idx = (hash.bits >> (52 - ((i + 1) * 5))) & 0x1f;
                }
                buf[i] = geoalphabet[idx];
            }
            buf[11] = '\0';
            addReplyBulkCBuffer(c, buf, 11);
        }
    }
}

Learn more:

目录

geo的结构

CRUD函数

geoadd命令

其他命令

geohash