索引和查询向量

了解如何在 Redis 中索引和查询向量嵌入

Redis 查询引擎允许您在 hash 或 JSON 对象中索引向量字段（有关更多信息，请参阅向量参考页面）。向量字段可以存储文本嵌入（text embeddings），文本嵌入是 AI 生成的向量表示，用来表示文本片段的语义信息。两个嵌入之间的向量距离表示它们在语义上的相似程度。通过比较由某些查询文本生成的嵌入与存储在 hash 或 JSON 字段中的嵌入之间的相似性，Redis 可以检索在含义上与查询密切匹配的文档。

在下面的示例中，我们使用 HuggingFace 模型 all-mpnet-base-v2 生成向量嵌入，并使用 Redis 查询引擎进行存储和索引。代码首先演示了 hash 文档，然后另起一节解释了与 JSON 文档的区别。

初始化

如果您使用的是 Maven，请将以下依赖项添加到您的 pom.xml 文件中

<dependency>
    <groupId>redis.clients</groupId>
    <artifactId>jedis</artifactId>
    <version>5.2.0</version>
</dependency>
<dependency>
    <groupId>ai.djl.huggingface</groupId>
    <artifactId>tokenizers</artifactId>
    <version>0.24.0</version>
</dependency>

如果您使用的是 Gradle，请将以下依赖项添加到您的 build.gradle 文件中

implementation 'redis.clients:jedis:5.2.0'
implementation 'ai.djl.huggingface:tokenizers:0.24.0'

导入依赖项

在您的源文件中导入以下类

// Jedis client and query engine classes.
import redis.clients.jedis.UnifiedJedis;
import redis.clients.jedis.search.*;
import redis.clients.jedis.search.schemafields.*;
import redis.clients.jedis.search.schemafields.VectorField.VectorAlgorithm;
import redis.clients.jedis.exceptions.JedisDataException;

// Data manipulation.
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.Map;
import java.util.List;
import org.json.JSONObject;

// Tokenizer to generate the vector embeddings.
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;

定义一个辅助方法

我们的嵌入模型将向量表示为长整型数组，但 Redis 查询引擎期望向量分量是浮点型值。此外，当您将向量存储在 hash 对象中时，必须将向量数组编码为字节字符串。为了简化这种情况，我们声明了一个辅助方法 longsToFloatsByteString()，该方法接受嵌入模型返回的长整型数组，将其转换为浮点型数组，然后将浮点型数组编码为字节字符串

public static byte[] longsToFloatsByteString(long[] input) {
    float[] floats = new float[input.length];
    for (int i = 0; i < input.length; i++) {
        floats[i] = input[i];
    }

    byte[] bytes = new byte[Float.BYTES * floats.length];
    ByteBuffer
        .wrap(bytes)
        .order(ByteOrder.LITTLE_ENDIAN)
        .asFloatBuffer()
        .put(floats);
    return bytes;
}

创建分词器实例

我们将使用 all-mpnet-base-v2 分词器生成嵌入。表示嵌入的向量具有 768 个分量，与输入文本的长度无关。

HuggingFaceTokenizer sentenceTokenizer = HuggingFaceTokenizer.newInstance(
    "sentence-transformers/all-mpnet-base-v2",
    Map.of("maxLength", "768",  "modelMaxLength", "768")
);

创建索引

连接到 Redis 并删除之前创建的名为 vector_idx 的索引。（如果索引不存在，ftDropIndex() 调用会抛出异常，这就是需要 try...catch 块的原因。）

UnifiedJedis jedis = new UnifiedJedis("redis://localhost:6379");

try {jedis.ftDropIndex("vector_idx");} catch (JedisDataException j){}

接下来，我们创建索引。下面示例中的 schema 包含三个字段：要索引的文本内容，一个tag字段表示文本的“体裁”，以及从原始文本内容生成的嵌入向量。 embedding 字段指定了HNSW索引、L2向量距离指标、表示向量分量的 Float32 值，以及 768 维度，这是 all-mpnet-base-v2 嵌入模型的要求。

FTCreateParams 对象指定使用 hash 对象进行存储，并使用 doc: 前缀来标识我们希望索引的 hash 对象。

SchemaField[] schema = {
    TextField.of("content"),
    TagField.of("genre"),
    VectorField.builder()
        .fieldName("embedding")
        .algorithm(VectorAlgorithm.HNSW)
        .attributes(
            Map.of(
                "TYPE", "FLOAT32",
                "DIM", 768,
                "DISTANCE_METRIC", "L2"
            )
        )
        .build()
};

jedis.ftCreate("vector_idx",
    FTCreateParams.createParams()
        .addPrefix("doc:")
        .on(IndexDataType.HASH),
        schema
);

添加数据

您现在可以提供数据对象，只要使用索引定义中指定的 doc: 前缀通过 hset() 添加它们，就会自动进行索引。

如下所示，使用 sentenceTokenizer 对象的 encode() 方法创建表示 content 字段的嵌入。紧随 encode() 之后的 getIds() 方法获取长整型向量，然后我们使用辅助方法将其转换为存储为字节字符串的浮点型数组。当您索引 hash 对象（如此处所示）时，请使用字节字符串表示；但对于 JSON 对象，请使用浮点型数组（参见下面的与 JSON 对象的区别）。请注意，当我们设置 embedding 字段时，必须使用 hset() 的重载方法，该方法要求 key、字段名和值都是字节数组，这就是我们在字符串上包含 getBytes() 调用的原因。

String sentence1 = "That is a very happy person";
jedis.hset("doc:1", Map.of("content", sentence1, "genre", "persons"));
jedis.hset(
    "doc:1".getBytes(),
    "embedding".getBytes(),
    longsToFloatsByteString(sentenceTokenizer.encode(sentence1).getIds())
);

String sentence2 = "That is a happy dog";
jedis.hset("doc:2", Map.of("content", sentence2, "genre", "pets"));
jedis.hset(
    "doc:2".getBytes(),
    "embedding".getBytes(),
    longsToFloatsByteString(sentenceTokenizer.encode(sentence2).getIds())
);

String sentence3 = "Today is a sunny day";
jedis.hset("doc:3", Map.of("content", sentence3, "genre", "weather"));
jedis.hset(
    "doc:3".getBytes(),
    "embedding".getBytes(),
    longsToFloatsByteString(sentenceTokenizer.encode(sentence3).getIds())
);

运行查询

创建索引并添加数据后，您就可以运行查询了。为此，您必须从您选择的查询文本创建另一个嵌入向量。Redis 在运行查询时会计算查询向量与索引中每个嵌入向量之间的向量距离。我们可以请求对结果进行排序，按升序距离排列。

下面的代码使用 encode() 方法创建查询嵌入，这与索引时相同，并在执行查询时将其作为参数传递（有关如何在查询中使用嵌入参数的更多信息，请参阅向量搜索）。该查询是一个K 近邻 (KNN)搜索，它按向量距离从查询向量的远近对结果进行排序。

String sentence = "That is a happy person";

int K = 3;
Query q = new Query("*=>[KNN $K @embedding $BLOB AS distance]")
                .returnFields("content", "distance")
                .addParam("K", K)
                .addParam(
                    "BLOB",
                    longsToFloatsByteString(
                        sentenceTokenizer.encode(sentence)..getIds()
                    )
                )
                .setSortBy("distance", true)
                .dialect(2);

List<Document> docs = jedis.ftSearch("vector_idx", q).getDocuments();

for (Document doc: docs) {
    System.out.println(
        String.format(
            "ID: %s, Distance: %s, Content: %s",
            doc.getId(),
            doc.get("distance"),
            doc.get("content")
        )
    );
}

假设您已将上述步骤中的代码添加到您的源文件中，现在可以运行了，但请注意，首次运行时可能需要一些时间才能完成（这是因为分词器在生成嵌入之前必须下载 all-mpnet-base-v2 模型数据）。运行代码时，会输出以下结果文本

Results:
ID: doc:2, Distance: 1411344, Content: That is a happy dog
ID: doc:1, Distance: 9301635, Content: That is a very happy person
ID: doc:3, Distance: 67178800, Content: Today is a sunny day

请注意，结果按 distance 字段的值排序，距离越小表示与查询的相似度越高。对于此模型，文本 "That is a happy dog" 被判定为在含义上与查询文本 "That is a happy person" 最相似的结果。

与 JSON 文档的区别

索引 JSON 文档与索引 hash 文档类似，但有一些重要的区别。JSON 允许使用嵌套字段进行更丰富的数据建模，因此您必须在 schema 中提供一个路径来标识要索引的每个字段。但是，您可以为每个路径声明一个简短的别名（使用 as() 选项），以避免在每次查询时都输入完整的路径。此外，在创建索引时，必须指定 IndexDataType.JSON。

下面的代码展示了这些区别，但索引在其他方面与之前为 hash 创建的索引非常相似

SchemaField[] jsonSchema = {
    TextField.of("$.content").as("content"),
    TagField.of("$.genre").as("genre"),
    VectorField.builder()
        .fieldName("$.embedding").as("embedding")
        .algorithm(VectorAlgorithm.HNSW)
        .attributes(
            Map.of(
                "TYPE", "FLOAT32",
                "DIM", 768,
                "DISTANCE_METRIC", "L2"
            )
        )
        .build()
};

jedis.ftCreate("vector_json_idx",
    FTCreateParams.createParams()
        .addPrefix("jdoc:")
        .on(IndexDataType.JSON),
        jsonSchema
);

JSON 索引的一个重要区别是使用浮点型数组而不是二进制字符串来指定向量。这需要对之前使用的 longsToFloatsByteString() 方法进行修改

public static float[] longArrayToFloatArray(long[] input) {
    float[] floats = new float[input.length];
    for (int i = 0; i < input.length; i++) {
        floats[i] = input[i];
    }
    return floats;
}

使用 jsonSet() 添加数据，而不是 hset()。使用 JSONObject 的实例提供数据，而不是像 hash 对象那样使用 Map。

String jSentence1 = "That is a very happy person";

JSONObject jdoc1 = new JSONObject()
        .put("content", jSentence1)
        .put("genre", "persons")
        .put(
            "embedding",
            longArrayToFloatArray(
                sentenceTokenizer.encode(jSentence1).getIds()
            )
        );

jedis.jsonSet("jdoc:1", Path2.ROOT_PATH, jdoc1);

String jSentence2 = "That is a happy dog";

JSONObject jdoc2 = new JSONObject()
        .put("content", jSentence2)
        .put("genre", "pets")
        .put(
            "embedding",
            longArrayToFloatArray(
                sentenceTokenizer.encode(jSentence2).getIds()
            )
        );

jedis.jsonSet("jdoc:2", Path2.ROOT_PATH, jdoc2);

String jSentence3 = "Today is a sunny day";

JSONObject jdoc3 = new JSONObject()
        .put("content", jSentence3)
        .put("genre", "weather")
        .put(
            "embedding",
            longArrayToFloatArray(
                sentenceTokenizer.encode(jSentence3).getIds()
            )
        );

jedis.jsonSet("jdoc:3", Path2.ROOT_PATH, jdoc3);

查询与 hash 文档的查询几乎相同。这表明为 JSON 路径选择合适的别名可以节省您编写复杂查询的工作。需要注意的一点是，查询的向量参数仍然指定为二进制字符串（使用 longsToFloatsByteString() 方法），即使 JSON 的 embedding 字段的数据被指定为数组。

String jSentence = "That is a happy person";

int jK = 3;
Query jq = new Query("*=>[KNN $K @embedding $BLOB AS distance]").
                    returnFields("content", "distance").
                    addParam("K", jK).
                    addParam(
                        "BLOB",
                        longsToFloatsByteString(
                            sentenceTokenizer.encode(jSentence).getIds()
                        )
                    )
                    .setSortBy("distance", true)
                    .dialect(2);

// Execute the query
List<Document> jDocs = jedis
        .ftSearch("vector_json_idx", jq)
        .getDocuments();

除了 key 的 jdoc: 前缀外，JSON 查询的结果与 hash 查询的结果相同

Results:
ID: jdoc:2, Distance: 1411344, Content: That is a happy dog
ID: jdoc:1, Distance: 9301635, Content: That is a very happy person
ID: jdoc:3, Distance: 67178800, Content: Today is a sunny day

了解更多

有关向量的索引选项、距离指标和查询格式的更多信息，请参阅向量搜索。

产品

工具

获取 Redis

连接

学习

最新

查看工作原理