100天跟着CP学PostgreSQL+AI,第14天 : 智能查询优化器：PostgreSQL 如何用 AI 生成执行计划?

作者介绍：崔鹏，计算机学博士，专注 AI 与大数据管理领域研究，拥有十五年数据库、操作系统及存储领域实战经验，兼具 ORACLE OCM、MySQL OCP 等国际权威认证，PostgreSQL ACE，运营技术公众号 "CP 的 PostgreSQL 厨房"，持续输出数据库技术洞察与实践经验。作为全球领先专网通信公司核心技术专家，深耕数据库高可用、高性能架构设计，创新探索 AI 在数据库领域的应用落地，其技术方案有效提升企业级数据库系统稳定性与智能化水平。学术层面，已在AI方向发表2篇SCI论文,将理论研究与工程实践深度结合，形成独特的技术研发视角。

系列文章介绍

第三阶段 : 技术深水区 AI 驱动数据库优化

主要内容

主题：智能查询优化器：PostgreSQL 如何用 AI 生成执行计划？

核心内容：传统 Cost-Based 优化 vs 机器学习优化 / 训练数据收集（pg_stat_statements）

实践案例：用历史查询日志训练一个轻量级优化模型（TensorFlow Lite）

正文

在数据库领域，查询优化是提升数据库性能的关键环节。PostgreSQL 作为一款强大的开源关系型数据库，一直在不断探索更高效的查询优化方法。近年来，随着人工智能技术的发展，将 AI 应用于查询优化成为了一个热门方向。本文将深入探讨 PostgreSQL 如何利用 AI 生成执行计划，对比传统 Cost - Based 优化与机器学习优化，介绍训练数据收集方法，并通过实践案例展示如何用历史查询日志训练一个轻量级优化模型（TensorFlow Lite）。

一、传统 Cost - Based 优化 vs 机器学习优化

（一）传统 Cost - Based 优化

传统的 PostgreSQL 查询优化器采用基于代价的优化（Cost - Based Optimization, CBO）方法。它通过估算不同执行计划的代价，选择代价最小的执行计划来执行查询。CBO 的核心是基于统计信息，如表的行数、列的基数、索引的使用情况等，来计算每个操作符（如连接、排序、过滤等）的代价。

CBO 的优点是具有明确的数学模型和理论基础，能够在一定程度上保证优化的合理性。然而，它也存在一些局限性。首先，统计信息的准确性对优化结果有很大影响。如果统计信息过时或不准确，优化器可能会选择错误的执行计划。其次，CBO 依赖于人工设计的代价函数，这些代价函数可能无法完全捕捉到实际执行环境中的各种因素，如硬件性能、数据分布的复杂性等。此外，对于一些复杂的查询，尤其是涉及多个表连接和复杂谓词的查询，CBO 的优化过程可能会非常耗时。

（二）机器学习优化

随着机器学习技术的发展，越来越多的研究开始将机器学习应用于数据库查询优化。机器学习优化的基本思想是利用历史查询数据和执行计划数据，训练一个模型来预测最优的执行计划。与传统的 CBO 相比，机器学习优化具有以下优势：

自动学习复杂模式：机器学习模型可以从大量的历史数据中自动学习数据分布、查询模式和执行计划之间的关系，无需人工设计复杂的代价函数。

适应动态环境：能够根据新的数据和查询动态调整优化策略，更好地适应数据库环境的变化。

处理复杂查询：对于复杂的查询，机器学习模型可以通过学习历史案例，找到更优的执行计划，提高查询执行效率。

二、训练数据收集（pg_stat_statements）

要训练一个机器学习模型来生成执行计划，首先需要收集大量的历史查询数据和对应的执行计划数据。在 PostgreSQL 中，可以使用 pg_stat_statements 扩展来收集这些数据。

（一）安装和启用 pg_stat_statements

确保 PostgreSQL 已经安装了 pg_stat_statements 扩展。在大多数情况下，pg_stat_statements 已经包含在 PostgreSQL 的安装包中。如果没有，可以通过以下命令安装（以 Ubuntu 为例）：

sudo apt - get install postgresql - contrib

在 PostgreSQL 的配置文件 postgresql.conf 中添加以下行来启用 pg_stat_statements：

shared_preload_libraries = 'pg_stat_statements'
pg_stat_statements.track = all

重新启动 PostgreSQL 服务，使配置生效。

（二）收集查询数据

启用 pg_stat_statements 后，它会记录每个执行的查询的详细信息，包括查询文本、执行次数、总执行时间、计划时间、执行的行数量等。可以通过以下 SQL 语句查询这些数据：

SELECT queryid, query, calls, total_time, plan_time, total_exec_time, rows
FROM pg_stat_statements;

其中，query 是查询文本，plan_time 是生成执行计划的时间，total_time 是查询的总执行时间，rows 是查询返回的行数。这些数据将作为训练机器学习模型的输入特征和目标变量。

三、实践案例：用历史查询日志训练一个轻量级优化模型（TensorFlow Lite）

（一）实验环境

操作系统：Ubuntu 20.04

PostgreSQL 版本：14

Python 版本：3.8

机器学习框架：TensorFlow Lite 2.6.0

（二）数据预处理

从 pg_stat_statements 中导出历史查询数据后，需要对数据进行预处理，以便用于训练机器学习模型。预处理步骤包括：

清洗数据：去除重复的查询、无效的查询（如语法错误的查询）和执行时间过长或过短的异常数据。

提取特征：将查询文本转换为机器学习模型可以处理的特征向量。可以提取的特征包括查询中的表名、列名、操作符（如 SELECT、FROM、WHERE、JOIN 等）、谓词条件、连接类型（如内连接、外连接等）、排序和分组信息等。此外，还可以将统计信息作为特征，如表的行数、列的基数等。

标注目标变量：目标变量是最优的执行计划。由于执行计划比较复杂，可以将其转换为一种简化的表示形式，如执行计划的代价或执行时间。在本实验中，我们以查询的实际执行时间作为目标变量，希望模型能够预测不同执行计划的执行时间，从而选择执行时间最短的计划。

以下是一个简单的数据预处理示例代码：

import pandas as pd
import re
# 读取从 pg_stat_statements 导出的数据
data = pd.read_csv('query_data.csv')
# 清洗数据：去除查询文本为空的行
data = data[data['query'].notna()]
# 提取查询中的表名
def extract_tables(query):
    tables = re.findall(r'\bFROM\s+(\w+)|JOIN\s+(\w+)', query, re.IGNORECASE)
    tables = [table for group in tables for table in group if table]
    return tables
data['tables'] = data['query'].apply(extract_tables)
# 提取查询中的操作符
def extract_operators(query):
    operators = re.findall(r'\b(SELECT|FROM|WHERE|JOIN|ORDER BY|GROUP BY)\b', query, re.IGNORECASE)
    return operators
data['operators'] = data['query'].apply(extract_operators)
# 提取统计信息（假设已经从 PostgreSQL 中获取了表的行数等统计信息）
# 这里只是一个示例，实际应用中需要根据具体情况获取统计信息
data['table_rows'] = [1000, 2000, 1500, ...]  # 假设的表行数数据
# 选择特征和目标变量
X = data[['table_rows', ...]]  # 特征向量
y = data['total_time']  # 目标变量（执行时间）

（三）模型构建

由于我们需要训练一个轻量级的优化模型，适合在 PostgreSQL 数据库环境中部署，因此选择 TensorFlow Lite 作为机器学习框架。TensorFlow Lite 提供了轻量级的模型表示和推理引擎，能够在资源受限的环境中高效运行。

我们构建一个简单的神经网络模型作为优化模型，该模型包含输入层、隐藏层和输出层。输入层的维度与特征向量的维度相同，隐藏层可以使用全连接层，输出层的维度为 1，用于预测查询的执行时间。

以下是模型构建的代码示例：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
# 构建模型
model = Sequential([
    Dense(64, activation='relu', input_shape=(X.shape[1],)),
    Dense(32, activation='relu'),
    Dense(1)
])
# 编译模型
model.compile(optimizer='adam', loss='mean_squared_error')
（四）模型训练
使用预处理后的数据对模型进行训练。训练过程中，采用随机梯度下降优化算法，通过最小化均方误差损失函数来调整模型的参数。
# 划分训练集和测试集
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 训练模型
history = model.fit(X_train, y_train, epochs=50, batch_size=32, validation_split=0.1)
（五）模型评估
训练完成后，需要对模型进行评估，以检验其性能。可以使用测试集数据计算模型的预测误差，如均方根误差（RMSE）、平均绝对误差（MAE）等。
# 评估模型
test_loss = model.evaluate(X_test, y_test)
print('Test Loss:', test_loss)

（六）模型部署

将训练好的模型转换为 TensorFlow Lite 格式，并部署到 PostgreSQL 数据库中。在 PostgreSQL 中，可以通过自定义函数调用 TensorFlow Lite 模型，实现基于 AI 的执行计划生成。

以下是将模型转换为 TensorFlow Lite 格式的代码示例：

# 转换为 TensorFlow Lite 模型
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
# 保存模型
with open('optimizer_model.tflite', 'wb') as f:
    f.write(tflite_model)

四、总结与展望

本文介绍了 PostgreSQL 如何利用 AI 生成执行计划，对比了传统 Cost - Based 优化与机器学习优化的优缺点，详细说明了训练数据收集方法（pg_stat_statements），并通过实践案例展示了如何用历史查询日志训练一个轻量级优化模型（TensorFlow Lite）。

通过将 AI 技术应用于查询优化，PostgreSQL 有望在复杂查询处理、动态环境适应等方面取得更好的性能。然而，目前的研究还处于初级阶段，仍然存在一些挑战，如如何提高模型的泛化能力、如何处理大规模数据、如何与传统优化器更好地结合等。

未来，随着机器学习技术的不断发展和数据库领域的深入研究，AI 在数据库查询优化中的应用将会更加广泛和成熟。我们可以期待看到更智能、更高效的查询优化器，为数据库性能提升带来新的突破。

如果你对 PostgreSQL 的智能查询优化器感兴趣，不妨尝试按照本文的方法进行实验，探索 AI 在数据库领域的更多可能性。