我们在 Pinterest 如何构建 Text-to-SQL

Adam Obeng | 数据科学家，数据平台科学；J.C. Zhong | 技术负责人，分析平台；Charlie Gu | 高级经理，工程部门

编写查询以解决分析问题是 Pinterest 数据用户的核 心任务。然而，在一个数据量庞大且分布在不同领域 的快速变化环境中，找到正确的数据并将分析问题转 换为正确且高效的 SQL 代码可能会是一项具有挑战 性的任务。

我们利用大型语言模型（LLMs）的普及作为机会，探索是否可以通过开发一个将这些分析问题直接转化为代码的Text-to-SQL功能来帮助我们的数据用户完成任务。

大多数在 Pinterest 上的数据分析都是通过我们的内部开源大数据 SQL 查询工具 Querybook 进行的。这个工具是我们开发和部署新功能以帮助数据用户（包括 Text-to-SQL）的自然选择。

第一个版本集成了一个简单的Text-to-SQL解决方案，利用了大型语言模型（LLM）。让我们更仔细地看一下它的架构：

用户提出一个分析性的问题，选择要使用的表。

1. 从表元数据存储中检索相关的表结构。
2. 将问题、选定的SQL方言和表结构编译成一个Text-to-SQL提示。
3. 将提示输入到大型语言模型（LLM）中。
4. 生成并显示给用户的流式响应。

表结构

从元数据存储中获取的表结构包括：

• 表名
• 表描述
• 列
• 列名
• 列类型
• 列描述

低基数列

某些分析查询，例如“‘web’平台上有多少活跃用户”，如果直接生成可能会产生不符合数据库实际值的SQL查询。例如，响应中的where子句可能是where platform='web'，而不是正确的where platform='WEB'。为了解决这类问题，会处理并整合低基数列的唯一值到表模式中，这样大语言模型可以利用这些信息生成精确的SQL查询。

上下文窗口限制

极大型的表结构可能会超过典型的上下文窗口限制。为了解决这个问题，我们采用了一些技术：

• 表结构简化版本：仅包含表名、列名和类型等关键元素。
• 列裁剪：在元数据存储中对列进行标记，并根据其标签从表结构中排除某些列。

一个大型语言模型（LLM）的完整响应可能需要数十秒，为了避免用户等待，我们采用了WebSocket来流式传输响应。由于需要返回生成的SQL之外的各种信息，因此一个结构良好的响应格式至关重要。虽然纯文本流式传输很简单，但流式传输JSON可能更复杂。我们采用了Langchain的部分JSON解析功能来进行服务器端的流式传输，然后将解析后的JSON通过WebSocket发送回客户端。

这里是我们目前用于Text2SQL的prompt：

评估与收获

我们对 Text-to-SQL 性能的初步评估主要是为了确保我们的实现与文献中报告的结果具有可比的性能，因为我们的实现主要采用了现成的方法。我们发现与文献中其他地方报告的结果 在 Spider 数据集上 相比，结果是可比的。尽管我们注意到该基准中的任务比我们用户面临的问题要简单得多，特别是它考虑的是一些预定义的表格，这些表格的列很少且标注得很好。

一旦我们的Text-to-SQL解决方案投入生产，我们也能够观察到用户如何与系统互动。随着我们的实现不断改进，用户也逐渐熟悉了这项功能，我们生成的SQL的一次性接受率从20%提高到了40%以上。在实践中，大多数生成的查询需要经过多次的人工或AI生成迭代才能最终完成。为了确定Text-to-SQL如何影响数据用户的生产力，最可靠的方法是进行实验。通过这种方法，先前的研究发现，AI辅助将任务完成速度提高了超过50%。在我们的实际数据（重要的是这些数据没有控制任务之间的差异）中，我们发现使用AI辅助编写SQL查询的任务完成速度提高了35%。

尽管第一个版本表现尚可——假设用户知道要使用的表——但在我们的数据仓库中从成千上万的表中识别出正确的表实际上对用户来说是一个巨大的挑战。为了解决这个问题，我们集成了检索增强生成（RAG），以引导用户选择正确的表来完成他们的任务。以下是对集成RAG的改进基础设施的回顾：

1. 使用离线任务生成表摘要和历史查询的向量索引。
2. 如果用户没有指定任何表，他们的问题将被转换为嵌入，并对向量索引进行相似性搜索，以推断出前 N 个合适的表。
3. 前 N 个表，连同表结构和分析问题，被编译成一个提示，供大语言模型（LLM）选择前 K 个最相关的表。
4. 前 K 个表被返回给用户进行验证或修改。
5. 使用用户确认的表继续标准的Text-to-SQL过程。

在向量索引中有两种类型的文档嵌入：

• 表摘要
• 查询摘要

Pinterest 正在进行一项持续的表标准化工作，为表添加分级。我们仅索引顶级表，促进使用这些高质量的数据集。表摘要生成过程包括以下步骤：

1. 从表元数据存储中检索表结构。
2. 收集使用该表的最新样本查询。
3. 根据上下文窗口，将尽可能多的样本查询与表结构一起纳入表摘要提示中。
4. 将提示发送给大语言模型以生成摘要。
5. 生成并存储向量存储中的嵌入。

表格摘要包括对表格的描述、表格包含的数据，以及潜在的使用场景。这是我们目前用于表格摘要的prompt：

除了在表摘要中的作用外，与每个表相关的示例查询也会单独进行摘要，包括查询的目的和所用表的详细信息。这里是我们使用的提示：

当用户提出一个分析性问题时，我们会使用相同的嵌入模型将其转换为嵌入。然后，我们对表和查询向量索引进行搜索。我们使用 OpenSearch 作为向量存储，并利用其内置的相似性搜索功能。

考虑到一个查询可能与多个表相关联，因此在相似度搜索结果中，同一个表可能会出现多次。目前，我们采用了一种简化的策略来聚合和评分这些表。表摘要的权重比查询摘要的权重更高，这种评分策略在未来可能会进行调整。

除了在 Text-to-SQL 中使用外，这种基于 NLP 的表搜索还用于 Querybook 中的一般表搜索。

在从向量索引中检索出前 N 个表之后，我们会使用一个大型语言模型（LLM）来选择与问题最相关的 K 个表，通过评估问题和表摘要来完成。根据上下文窗口的大小，我们会尽可能多地将表包含在提示中。这是我们在表重新选择时使用的提示：提示。

一旦表格被重新选择，它们会被返回给用户进行验证，然后再过渡到实际的SQL生成阶段。

我们通过之前表搜索的离线数据评估了我们Text-to-SQL功能中的表检索组件。这些数据在一个重要方面是不足的：它捕捉的是用户在知道基于NLP的搜索可用之前的行为。因此，这些数据主要用于确保基于嵌入的表搜索不会比现有的基于文本的搜索表现更差，而不是试图衡量改进。我们使用此评估来选择用于表检索的嵌入方法和设置权重。这种方法揭示了我们通过数据治理工作生成的表元数据对整体性能具有重要意义：没有表文档的嵌入搜索命中率为40%，但随着对表文档权重的增加，性能线性增长，最高可达90%。

虽然我们目前实现的Text-to-SQL已经显著提升了数据分析师的生产力，但仍存在改进的空间。以下是一些可能的进一步发展方向：

• 元数据增强

目前，我们的向量索引仅与表摘要关联。一个潜在的改进可能是包含更多元数据，例如分级、标签、领域等，以便在检索相似表时进行更精细的过滤。

• 定时或实时索引更新

当前向量索引是手动生成的。实现定期或甚至实时更新，每当有新表创建或查询执行时，可以提高系统效率。

• 相似性搜索和评分策略修订

我们当前用于聚合相似性搜索结果的评分策略相当基础。通过微调这一方面，可以提高检索结果的相关性。

目前，LLM 生成的 SQL 查询直接返回给用户而不进行验证，这留下了查询可能无法按预期运行的风险。通过实现查询验证，例如使用受限的 beam 搜索，可以提供额外的保障。

引入一个用户界面，以高效地收集用户对表格搜索和查询生成结果的反馈，可以为改进提供宝贵的见解。这些反馈可以被处理并融入到向量索引或表格元数据存储中，最终提升系统性能。

在进行这个项目的过程中，我们发现实际环境中文本到SQL的性能与现有基准测试中的表现存在显著差异。现有的基准测试倾向于使用少量预定义且高度规范化的表格。对于应用研究人员来说，生成更贴近现实的基准测试将非常有帮助，这些基准测试应包含更多非规范化的表格，并将表格搜索作为问题的核心部分。

要了解更多关于Pinterest工程的信息，请查看我们的 Engineering Blog ，并访问我们的 Pinterest Labs网站。要浏览和申请开放职位，请访问我们的 Careers 页面。