【kv存储】结合faiss构建向量内存数据库

一、介绍

Faiss是一个由Meta（原Facebook）人工智能实验室（FAIR）开发的高效相似性搜索与稠密向量聚类库，支持CPU和GPU两种计算方式。

KV 存储是一种内存数据库，准确来说，是一种精确存储的内存数据库，即key-value一一对应。我们在某些AI应用中（如图片识别、自然语言处理、推荐系统等），需要支持通过向量相似度来进行查询。为了支持这种AI应用，本文介绍如何结合FAISS为 KV 存储设计向量功能。

补充知识：向量化是一切的基础：在Faiss的世界里，任何数据（如一张图片、一段文本）都会被转换为一个"特征向量"，即一串数字。相似性搜索就变成了在数学空间中寻找最邻近向量的计算问题。

二、架构设计

我们先介绍下向量功能的架构设计：我对外设计了VSET、VGET、VDEL三个交互命令（V代表vector，命名类似于Redis的SET、GET、DEL命令）：

VSET vector key value
VGET vector threshold
VDEL key

功能如下：

VSET ：向faiss index里插入一条向量，并分配一个ID，在id-key映射表里面插入ID和key，在普通kv数据引擎（哈希表/红黑树/跳表）中插入key-value键值对

VGET ：遍历计算faiss index中向量与所查询的向量的相似度，在faiss index里寻找相似度最高的向量，比较阈值，超过阈值则取出ID，通过ID寻找key、value，最终返回value。

VDEL：根据key删除普通kv数据引擎中的key-value键值对，并删除其他两个表中的ID-value，id-vector映射关系。

了解了基本的架构之后，我们介绍简单的向量实现：

三、实现

faiss 提供了 faiss index 的 C 接口，faiss index里面存了 ID 和 vector

KV存储需要维护一个id-key映射表，可以选择哈希表来实现，用于 key 和 ID 的映射，faiss index 认 ID（内部自增的ID）。

这里需要多说明下关于这个ID的问题，IndexFlatL2的基础索引拥有自己的一套自增custom_ID，并不能天然使用外部定义的ID。解决的办法是用IndexIDMap来将flat_index包装成支持ID的索引。

c 复制代码

//创建基础索引
FaissIndex *flat_index;
faiss_IndexFlatL2_new(&flat_index, dimension);
//装饰器包装
FaissIndex *id_map_index;
// 将 flat_index 包装成支持 ID 的索引
faiss_IndexIDMap_new(&id_map_index, flat_index);

对外给出初始化/销毁函数和三个分发入口：

初始化：kvs_vector_init()创建faiss专用的id-key哈希表等，kvs_vector_destroy()销毁哈希表、free索引

c 复制代码

int kvs_vector_init(void) {
    memset(&g_vec, 0, sizeof(g_vec));
	//创建哈希表
    if (kvs_hash_create(&g_vec.id2key) != 0)
        return -1;
    if (kvs_hash_create(&g_vec.key2id) != 0) {
        kvs_hash_destory(&g_vec.id2key);
        memset(&g_vec, 0, sizeof(g_vec));
        return -1;
    }
    if (kvs_hash_create(&g_vec.store) != 0) {
        kvs_hash_destory(&g_vec.key2id);
        kvs_hash_destory(&g_vec.id2key);
        memset(&g_vec, 0, sizeof(g_vec));
        return -1;
    }
    g_vec.hashes_ready = 1;
    g_vec.next_id = 0;
    return 0;
}

void kvs_vector_destroy(void) {
    if (g_vec.index) {
        faiss_Index_free(g_vec.index); //释放faiss index
        g_vec.index = NULL;
    }
    //销毁哈希表
    if (g_vec.hashes_ready) {
        kvs_hash_destory(&g_vec.store);
        kvs_hash_destory(&g_vec.key2id);
        kvs_hash_destory(&g_vec.id2key);
    }
    memset(&g_vec, 0, sizeof(g_vec));
}

VSET入口函数：kvs_vector_set() -> faiss_Index_add_with_ids()，值得注意的是，我们的index是懒创建的，即第一次VSET写入的时候，确定维度并创建index

c 复制代码

int kvs_vector_set(const char *blob, size_t blob_n,
                   const char *key, size_t klen,
                   const char *val, size_t vlen) {
    int dim;
    idx_t id;
    float *vec;
    FaissIndexFlatL2 *flat = NULL;
    FaissIndexIDMap *id_map = NULL;

    dim = blob_dim(blob_n);
    vec = copy_blob(blob, dim);

    // 懒创建索引：如果不存在则创建 FlatL2 + IDMap
    if (!g_vec.index) {
        faiss_IndexFlatL2_new_with(&flat, (idx_t)dim);
        faiss_IndexIDMap_new(&id_map, (FaissIndex *)flat);
        faiss_IndexIDMap_set_own_fields(id_map, 1);
        g_vec.index = (FaissIndex *)id_map;
        g_vec.dim = dim;
    }

    id = g_vec.next_id;
    faiss_Index_add_with_ids(g_vec.index, 1, vec, &id);
    kvs_free(vec);

    kvs_hash_set(&g_vec.key2id, key, klen, (const char *)&id, sizeof(id));
    kvs_hash_set(&g_vec.id2key, (const char *)&id, sizeof(id), key, klen);
    kvs_hash_set(&g_vec.store, key, klen, val, vlen);

    g_vec.next_id++;
    return 0;
}

VGET入口函数：kvs_vector_get(vector_blob, threshold) -> faiss_Index_search()

c 复制代码

char *kvs_vector_get(const char *blob, size_t blob_n, float threshold,
                     size_t *vlen) {
    int dim = blob_dim(blob_n);
    float *vec = copy_blob(blob, dim);
    
    idx_t labels[1];
    float distances[1];
    
    faiss_Index_search(g_vec.index, 1, vec, 1, distances, labels); // 搜索faiss index中的近似向量
    free(vec);
    //通过id寻找key、value
    idx_t id = labels[0];
    size_t klen;
    char *key = kvs_hash_get(&g_vec.id2key, (char*)&id, sizeof(id), &klen);
    
    return kvs_hash_get(&g_vec.store, key, klen, vlen);
}

VDEL入口函数：kvs_vector_delete(key) -> 在普通哈希表中通过user_key查ID -> faiss_Index_remove() ->清理其他映射条目

c 复制代码

int kvs_vector_delete(const char *key, size_t klen) {
    char *id_buf;
    size_t id_len = 0;
    idx_t id;

    if (!key || !g_vec.hashes_ready)
        return -1;

    id_buf = kvs_hash_get(&g_vec.key2id, key, klen, &id_len);
    if (!id_buf || id_len != sizeof(idx_t))
        return 1;

    memcpy(&id, id_buf, sizeof(id));

    kvs_hash_del(&g_vec.store, key, klen);
    kvs_hash_del(&g_vec.id2key, (const char *)&id, sizeof(id));
    kvs_hash_del(&g_vec.key2id, key, klen);
    if (g_vec.index)
        remove_faiss_id(id);
    return 0;
}

总结：faiss提供了非常方便的操作index的接口函数，我们这里实现的是最简单的暴力搜寻计算L2，当数据规模较大的时候（通常大于一万），faiss提供了PQ、IVF、HNSW 等等高级的算法。faiss还提供了不同的应用类型下，不同的距离的计算方式如余弦相似度、内积等。