Spring AI 结合 PGVector 实现全量数据集向量化

起凡2026/3/30大约 5 分钟

在上一篇文章中，我们将 BIRD-SQL 数据集成功解析并持久化到了关系型数据库（PostgreSQL）中。然而，大模型（LLM）并不懂我们的关系型表结构，它需要通过“向量检索”来寻找与用户问题最相关的表、列和业务知识。

今天，我们将正式引入 Spring AI 和 PostgreSQL 的 pgvector 插件，把干巴巴的关系型数据转化为高维空间的“向量（Embedding）”，为我们后续构建 Agent 的 RAG（检索增强生成）能力打下坚实基础。

1. PostgreSQL 向量化基石：安装 pgvector 插件

要把 PostgreSQL 变成一个强大的向量数据库，我们需要安装 pgvector 插件。

initialize-schema: true 会让 Spring AI 自动为你创建一张名为 vector_store 的表。如果不填 dimensions 它其实也能运行，但如果你更换了底层的 Embedding 模型（比如从 1024 维换成了 1536 维的 OpenAI 模型），你需要手动删除数据库中的 vector_store 表并让系统重建，否则会报维度不匹配的致命错误！

3. Jimmer 关联查询优化：定义 Fetcher

在上一篇中我们定义过 DbColumn.toDocument() 扩展方法，该方法需要获取列所属表的 databaseId 和 id：

DataAgentSpec.Retrieval.DocumentMetadataKey.DATABASE_ID to dbTable.databaseId,
DataAgentSpec.Retrieval.DocumentMetadataKey.TABLE_ID to dbTable.id

如果直接使用普通的 findAll 或 where 查询出 DbColumn，在访问 dbTable 属性时会触发 N+1 查询甚至报错。借助 Jimmer 强大的 Fetcher 机制，我们可以精细化控制关联抓取：

interface DbColumnRepository : KRepository<DbColumn, UUID> {
    companion object {
        // 定义 Fetcher，一并抓取关联的 dbTable 所有标量字段
        val FETCHER = newFetcher(DbColumn::class).by {
            allScalarFields()
            dbTable {
                allScalarFields()
            }
        }
    }

    fun findByDatabaseId(databaseId: String): List<DbColumn> {
        return sql.createQuery(DbColumn::class) {
            where(table.dbTable.databaseId.eq(databaseId))
            select(table.fetch(FETCHER)) // 使用 Fetcher 查询
        }.execute()
    }
}

4. 全量数据向量化实战

万事俱备，我们可以编写测试用例 DatasetEmbeddingTest，将我们存入的表、列、外部知识全部转化为向量，送入 PGVector！

@SpringBootTest
class DatasetEmbeddingTest(
    @Autowired private val dbTableRepository: DbTableRepository,
    @Autowired private val dbColumnRepository: DbColumnRepository,
    @Autowired private val questionKnowledgeRepository: QuestionKnowledgeRepository,
    @Autowired private val glossaryKnowledgeRepository: GlossaryKnowledgeRepository,
    @Autowired private val vectorStore: VectorStore
) {
    @Test
    fun embeddingTest() {
        val databaseId = "california_schools"
        
        // 1. 获取所有实体，并将其转换为 Spring AI 的 Document 对象
        val documents = buildList {
            addAll(dbTableRepository.findByDatabaseId(databaseId).map { it.toDocument() })
            addAll(dbColumnRepository.findByDatabaseId(databaseId).map { it.toDocument() })
            addAll(questionKnowledgeRepository.findByDatabaseId(databaseId).map { it.toDocument() })
            addAll(glossaryKnowledgeRepository.findByDatabaseId(databaseId).map { it.toDocument() })
        }
        
        // 2. 分批次进行向量化，防止一次性请求过大导致 API 超时或触发限流
        documents.chunked(10).forEach { batchDocs ->
            vectorStore.add(batchDocs) 
            // 内部会调用 Embedding 模型获取向量，并写入 PG 的 vector_store 表
        }
    }

💡 代码设计亮点

我们使用了 Kotlin 的 chunked(10) 对文档进行了分块。由于大模型的 Embedding API 通常有并发和请求体大小（Token 限制），分批写入能极大提高入库的稳定性。

5. 验证向量检索 (Vector Retrieval)

存进去之后，我们肯定要测试一下“能不能精准搜出来”。由于我们将多种结构（表、列、知识库）都放在了同一个 vector_store 里，检索时必须借助 元数据过滤 (Metadata Filter) 来精准控制搜索范围。

下面的测试演示了如何精准搜索 california_schools 数据库下的 “表结构 (TABLE)” 信息：

    @Test
    fun retrieveTest() {
        val builder = FilterExpressionBuilder()
        
        // 构造过滤器：WHERE vectorType = 'table' AND databaseId = 'california_schools'
        val filterExpression = builder.and(
            builder.eq(
                DataAgentSpec.Retrieval.DocumentMetadataKey.VECTOR_TYPE,
                DataAgentSpec.Retrieval.VectorType.TABLE
            ),
            builder.eq(
                DataAgentSpec.Retrieval.DocumentMetadataKey.DATABASE_ID, 
                "california_schools"
            )
        ).build()

        // 构造搜索请求
        val request = SearchRequest.builder()
            .query("What is the highest eligible free rate for K-12 students in the schools in Alameda County?")
            .filterExpression(filterExpression) // 应用元数据过滤
            .topK(5) // 取最相似的前 5 条
            .build()
            
        // 执行相似度检索
        val documents = vectorStore.similaritySearch(request)
        
        // 按照相似度得分降序排序并打印
        documents.sortByDescending { it.score }
        logger.info { "检索结果：$documents" }
    }
}

6. 总结与预告

在本文中，我们不仅为 PostgreSQL 插上了 pgvector 这双翅膀，还利用 Spring AI 极简的 API 设计，完成了异构数据（表、列、业务问题）的统一建模、向量化与存储。更重要的是，我们验证了带元数据过滤的向量检索能力。

至此，我们的“知识库储备”已经全部就绪！在下一篇文章中，我们将踏入系列的核心高潮：借助 A2A 协议与 StateGraph（状态图），亲手编排一个能根据检索上下文自动生成高水平 SQL 的智能 Agent！ 敬请期待！