MybatisPlus + PostgreSQL 实现向量检索

YAYO老祖

已于 2025-03-17 21:46:27 修改

阅读量1.1k

点赞数 16

CC 4.0 BY-SA版权

分类专栏： postgresql 向量 mybatisplus 文章标签： postgresql 数据库 mybatis 向量

于 2025-03-17 10:33:07 首次发布

本文链接：https://blog.csdn.net/qq_29213799/article/details/146308124

postgresql 同时被 3 个专栏收录

3 篇文章

订阅专栏

向量

2 篇文章

订阅专栏

mybatisplus

1 篇文章

订阅专栏

1.前置条件

已经安装好postgresql，并且数据库安装了pgvector插件，可以创建带有向量字段的数据表

2.创建表带有向量字段vector， vector类型维度1024

CREATE TABLE document_chunk
(
    id          VARCHAR(64) PRIMARY KEY,
    document_id VARCHAR(64),
    chunk       VARCHAR(1024),
    chunk_index INT,
    vector      vector(1024)
);

COMMENT ON TABLE document_chunk IS '文档分片表';
COMMENT ON COLUMN document_chunk.id IS '主键id';
COMMENT ON COLUMN document_chunk.document_id IS '所属文档id';
COMMENT ON COLUMN document_chunk.chunk IS '文档分片内容';
COMMENT ON COLUMN document_chunk.chunk_index IS '分片索引';
COMMENT ON COLUMN document_chunk.vector IS '向量';

3.postgresql vector类型的理解，存储

postgresql 的 vector类型本质上就是一个向量字符串

所以存储的时候存储的向量字符串 '[-0.0123, 0.0456, ...]'

1.需要自定义一个类型转换器来将我们 java实体中的 List<Double>（这里也可以写 float[]，看你喜欢使用哪种类型去存储向量）转为为 '[-0.0123, 0.0456, ...]' 这种向量字符串字符串

import org.apache.ibatis.type.BaseTypeHandler;
import org.apache.ibatis.type.JdbcType;
import org.postgresql.util.PGobject;

import java.sql.CallableStatement;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * 由于 MyBatis-Plus 不直接支持自定义的扩展类型（如 pgvector），你需要自定义一个 TypeHandler 来处理 vector 类型的序列化和反序列化。
 * TypeHandler 字段类型处理器：在 MyBatis 中，类型处理器（TypeHandler）扮演着 JavaType 与 JdbcType 之间转换的桥梁角色。它们用于在执行 SQL 语句时，将 Java 对象的值设置到 PreparedStatement 中，或者从 ResultSet 或 CallableStatement 中取出值。
 *
 **/
public class PGVectorTypeHandler extends BaseTypeHandler<List<Double>> {

    /**
     * 将 Java 对象的值设置到 PreparedStatement 中
     * java的值写到数据库中去
     */
    @Override
    public void setNonNullParameter(PreparedStatement ps, int i, List<Double> parameter, JdbcType jdbcType) throws SQLException {
        PGobject vector = new PGobject();
        vector.setType("vector");
        StringBuilder sb = new StringBuilder();
        sb.append("[");
        for (int j = 0; j < parameter.size(); j++) {
            if (j > 0) sb.append(",");
            sb.append(parameter.get(j));
        }
        sb.append("]");
        vector.setValue(sb.toString());
        ps.setObject(i, vector);
    }

    /**
     * 从 ResultSet 或 CallableStatement 中取出值
     * 从数据库中的值读取到java对象中
     */
    @Override
    public List<Double> getNullableResult(ResultSet rs, String columnName) throws SQLException {
        String vecStr = rs.getString(columnName);
        return parseVector(vecStr);
    }

    /**
     * 从 ResultSet 或 CallableStatement 中取出值
     * 从数据库中的值读取到java对象中
     */
    @Override
    public List<Double> getNullableResult(ResultSet rs, int columnIndex) throws SQLException {
        String vecStr = rs.getString(columnIndex);
        return parseVector(vecStr);
    }

    /**
     * 从 CallableStatement 中取出值
     * 从数据库中的值读取到java对象中
     */
    @Override
    public List<Double> getNullableResult(CallableStatement cs, int columnIndex) throws SQLException {
        String vecStr = cs.getString(columnIndex);
        return parseVector(vecStr);
    }

    /**
     * 将值设置到 PreparedStatement 中
     */
    private List<Double> parseVector(String vecStr) {
        if (vecStr == null || vecStr.isEmpty()) return null;
        // 去掉字符串首尾的 [ ] 只获取中间数字
        vecStr = vecStr.substring(1, vecStr.length() - 1);
        String[] parts = vecStr.split(",");
        List<Double> vector = new ArrayList<>(parts.length);
        for (int i = 0; i < parts.length; i++) {
            vector.add(Double.parseDouble(parts[i].trim()));
        }
        return vector;
    }
}

2.Entity实体类中使用上述类型转换器

import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.annotation.TableId;
import com.baomidou.mybatisplus.annotation.TableName;

import java.io.Serializable;
import java.util.List;

import xxx.xxx.xxx.PGVectorTypeHandler;
import lombok.Data;
import lombok.EqualsAndHashCode;


@Data
@EqualsAndHashCode(callSuper = false)
@TableName("document_chunk")
public class DocumentChunkEntity implements Serializable {

    private static final long serialVersionUID = 1L;

    /**
     * 主键id
     */
    @TableId("id")
    private String id;

    /**
     * 所属文档id
     */
    @TableField("document_id")
    private String documentId;

    /**
     * 文档分片内容
     */
    @TableField("chunk")
    private String chunk;

    /**
     * 分片索引
     */
    @TableField("chunk_index")
    private Integer chunkIndex;

    /**
     * 向量
     * <p>
     * typeHandler 配置 字段处理器
     */
    @TableField(value = "vector", typeHandler = PGVectorTypeHandler.class)
    private List<Double> vector;


}

3.sql xml 文件中使用上述类型解析器

<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
<mapper namespace="xxx.xxx.xxx.DocumentChunkMapper">

    <!-- 通用查询映射结果 -->
    <resultMap id="BaseResultMap" type="xxx.xxx.xxx.DocumentChunkEntity">
        <id column="id" property="id"/>
        <result column="document_id" property="documentId"/>
        <result column="chunk" property="chunk"/>
        <result column="chunk_index" property="chunkIndex"/>
        <!-- XML 配置 字段处理器 -->
        <result column="vector" property="vector" typeHandler="xxx.xxx.xxx.PGVectorTypeHandler"/>
    </resultMap>

    <!-- 通用查询结果列 -->
    <sql id="Base_Column_List">
        id, document_id, chunk, chunk_index, vector
    </sql>

    <!--
        1.欧几里得距离:欧几里得距离是一种用于度量两点之间“直线距离”的方式。以下查询会返回与给定向量 [3, 3, 3] 欧几里得距离最小的记录：
        SELECT id, embedding FROM items ORDER BY embedding <-> '[3, 3, 3]' LIMIT 1;
        <-> 运算符用于计算向量之间的欧几里得距离，并按照从小到大的顺序排序结果。

        2. 余弦相似度
        余弦相似度用于衡量两个向量的夹角，余弦相似度越接近 1，表示两个向量越相似。以下查询会返回与给定向量 [3, 3, 3] 余弦相似度最高的记录：
        SELECT id, embedding FROM items ORDER BY embedding <=> '[3, 3, 3]' LIMIT 1;
        <=> 运算符用于计算余弦相似度，并按降序排列结果。

        3. 内积
        内积（Dot Product）用于测量两个向量在相同方向上的相似性，常用于推荐系统。以下查询会返回与给定向量 [3, 3, 3] 内积最大的记录
        SELECT id, embedding FROM items ORDER BY embedding <#> '[3, 3, 3]' LIMIT 1;

        错误：operator does not exist: vector <=> real[]
        原因：未正确将向量转换为 VECTOR 类型。
        解决方案：
        确保检索时传入的向量是字符串格式（如 '[1,2,3]'）。
        在 SQL 中使用 CAST(#{vector} AS vector) 或 #{vector}::vector 显式转换类型。
    -->

    <!-- 使用pgvector的 -> operator 进行向量搜索 -->
    <select id="searchByVector" resultMap="BaseResultMap">
        SELECT
        <include refid="Base_Column_List"/>
        FROM document_chunk
        ORDER BY vector <![CDATA[<=>]]> #{queryVector}::vector
        LIMIT 5
    </select>
</mapper>

4.mapper 中添加检索方法

import xxx.xxx.xxx.DocumentChunkEntity;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;

import java.util.List;

/**
 * <p>
 * 文档分片表 Mapper 接口
 * </p>
 *
 */
@Mapper
public interface DocumentChunkMapper extends BaseMapper<DocumentChunkEntity> {

    /**
     * 向量余弦相似度搜索
     * @param queryVector
     * @return
     */
    List<DocumentChunkEntity> searchByVector(@Param("queryVector") String queryVector);

}

4.插入向量数据

这边使用bge-m3向量模型对文档的片段进行向量处理，然后存储到数据库中

import com.plexpt.chatgpt.ChatGPT;
import com.plexpt.chatgpt.entity.embedding.EmbeddingData;
import com.plexpt.chatgpt.entity.embedding.EmbeddingRequest;
import com.plexpt.chatgpt.entity.embedding.EmbeddingResult;

private final ChatGPT chatGPT = ChatGPT.builder()
		.apiHost("http://localhost:11434/")
		.apiKeyList(ListUtils.of("sk-xxx"))
		.build()
		.init();



EmbeddingRequest embeddingRequest = EmbeddingRequest.builder()
		.model("bge-m3")
		.input(chunks)
		.build();
EmbeddingResult embeddingResult = chatGPT.createEmbeddings(embeddingRequest);
List<EmbeddingData> embeddings = embeddingResult.getData();

List<DocumentChunkEntity> documentChunkEntityList = new ArrayList<>();
for (int i = 0; i < chunks.size(); i++) {
	DocumentChunkEntity documentChunk = new DocumentChunkEntity();
	documentChunk.setId(IdUtils.getSnowflakeNextIdStr());
	documentChunk.setChunk(chunks.get(i));
	documentChunk.setDocumentId(document.getId());
	documentChunk.setChunkIndex(i);
	// 向量 List<Double>
	documentChunk.setVector(embeddings.get(i).getEmbedding());
	documentChunkEntityList.add(documentChunk);
}
documentChunkService.saveBatch(documentChunkEntityList);

5.向量检索

一样使用bge-m3对输入进行向量化，然后检索出近似最高的几个chunk

import com.plexpt.chatgpt.ChatGPT;
import com.plexpt.chatgpt.entity.embedding.EmbeddingData;
import com.plexpt.chatgpt.entity.embedding.EmbeddingRequest;
import com.plexpt.chatgpt.entity.embedding.EmbeddingResult;

private final ChatGPT chatGPT = ChatGPT.builder()
		.apiHost("http://localhost:11434/")
		.apiKeyList(ListUtils.of("sk-xxx"))
		.build()
		.init();



public List<String> search(String prompt) {
	// embedding 请求获取向量
	EmbeddingRequest embeddingRequest = EmbeddingRequest.builder()
			.input(ListUtils.of(prompt))
			.model("bge-m3")
			.build();
	List<Double> embeddings = chatGPT.createEmbeddings(embeddingRequest)
			.getData().get(0).getEmbedding();

	List<DocumentChunkEntity> chunkEntities = documentChunkMapper.searchByVector(embeddings.toString());

	log.info("chunkEntities:{} 数量", chunkEntities.size());

	return chunkEntities.stream().map(DocumentChunkEntity::getChunk).collect(Collectors.toList());
}