如何在Redis向量数据库中执行批量搜索？

直接传入向量列表的批量搜索方式目前不可行，Redis Stack的Search命令（以及redis-py的封装）本身不支持单次请求中传入多个向量执行多轮KNN搜索。以下是两种可行的替代实现方案：

方案1：用Pipeline批量执行搜索请求 #

借助redis-py的Pipeline功能，把多个搜索请求打包成批量操作，能有效减少网络往返开销，这是最直接的解决方案。

示例代码：

import redis
import numpy as np
from redis.commands.search.query import Query

VECTOR_DIMENSIONS = 128
INDEX_NAME = "my_vector_index"

client = redis.Redis(host="localhost", port=6379, decode_responses=True)

# 准备多个查询向量
query_vectors = [
    np.random.rand(VECTOR_DIMENSIONS).astype(np.float32).tobytes(),
    np.random.rand(VECTOR_DIMENSIONS).astype(np.float32).tobytes()
]

# 创建Pipeline对象
pipe = client.pipeline()

# 给每个向量添加对应的搜索命令
for vec in query_vectors:
    query = (
        Query("*=>[KNN 2 @vector $vec as score]")
        .sort_by("score")
        .return_fields("id", "score")
        .paging(0, 2)
        .dialect(2)
    )
    pipe.ft(INDEX_NAME).search(query, {"vec": vec})

# 执行批量请求并获取所有结果
results = pipe.execute()

# 遍历处理每个向量的搜索结果
for idx, result in enumerate(results):
    print(f"第{idx+1}个向量的搜索结果：")
    for doc in result.docs:
        print(f"id: {doc.id}, score: {doc.score}")

方案2：用Redis Lua脚本实现（进阶） #

如果追求更高的执行效率，可以编写Lua脚本在Redis服务器端直接执行多轮KNN搜索，进一步降低网络传输带来的延迟。

示例Lua脚本：

local index_name = KEYS[1]
local vector_list = ARGV
local results = {}

for i, vec in ipairs(vector_list) do
    -- 构建并执行FT.SEARCH命令
    local resp = redis.call("FT.SEARCH", index_name, "*=>[KNN 2 @vector $vec as score]",
        "SORTBY", "score",
        "RETURN", "2", "id", "score",
        "LIMIT", "0", "2",
        "DIALECT", "2",
        "PARAMS", "2", "vec", vec)
    table.insert(results, resp)
end

return results

Python调用示例：

# 注册并执行Lua脚本
lua_script = """
local index_name = KEYS[1]
local vector_list = ARGV
local results = {}

for i, vec in ipairs(vector_list) do
    local resp = redis.call("FT.SEARCH", index_name, "*=>[KNN 2 @vector $vec as score]",
        "SORTBY", "score",
        "RETURN", "2", "id", "score",
        "LIMIT", "0", "2",
        "DIALECT", "2",
        "PARAMS", "2", "vec", vec)
    table.insert(results, resp)
end

return results
"""

script = client.register_script(lua_script)
results = script(keys=[INDEX_NAME], args=query_vectors)

补充说明 #

• Pipeline是将多个请求打包发送，Redis会逐个执行，但能避免多次网络请求的延迟损耗。
• Lua脚本在Redis服务器端运行，性能更优，但要注意脚本的执行时长，避免影响Redis的正常服务。
• 当前Redis Stack（含RediSearch模块）暂不支持单次FT.SEARCH命令传入多向量执行批量KNN搜索，后续版本可能会新增相关功能，可关注官方更新。

内容的提问来源于stack exchange，提问作者gu1show