AI 驱动的查询优化：当执行计划遇上机器学习，数据库内核的智能进化

国医中兴

2019人浏览 · 2026-06-30 12:02:54

国医中兴 · 2026-06-30 12:02:54 发布

AI 驱动的查询优化：当执行计划遇上机器学习，数据库内核的智能进化

一、执行计划选错引发的性能灾难：传统优化器的天花板

在分布式存储与数据库内核的开发实践中，查询优化器始终是最核心也最脆弱的组件。传统基于代价的优化器（CBO）依赖统计信息进行计划选择，但当数据分布倾斜、多表关联复杂度上升时，统计信息的滞后性和估算误差会导致优化器做出灾难性的执行计划选择。

一个典型的生产场景：某电商平台的大促期间，订单表从日均 500 万行暴增至 3 亿行，而统计信息仍停留在 T+1 的采集周期。优化器基于过时的统计信息选择了 Nested Loop Join 而非 Hash Join，导致一条本应 200ms 完成的查询执行了 47 分钟，直接拖垮了整个数据库实例的连接池。

传统优化器的根本瓶颈在于：它只能基于静态模型做一次性决策，无法从历史执行反馈中学习。而 AI 驱动的查询优化正是要突破这一天花板——让优化器从"规则驱动"进化为"数据驱动"。

二、从代价模型到学习型优化器：AI 查询优化的底层机制

AI 驱动的查询优化并非简单地在优化器外层套一个神经网络，而是需要深入理解查询优化的决策链路，在关键节点嵌入学习与推理能力。

flowchart TB
    SQL[SQL 查询输入] --> Parser[语法解析器]
    Parser --> AST[抽象语法树]
    AST --> Rewriter[逻辑改写器]
    Rewriter --> LP[逻辑计划]
    LP --> CBO[传统代价优化器]
    LP --> ML_Model[学习型代价模型]
    CBO --> Plan_A[候选计划 A]
    CBO --> Plan_B[候选计划 B]
    CBO --> Plan_C[候选计划 C]
    ML_Model --> Cost_Pred[代价预测]
    Cost_Pred --> Selector[计划选择器]
    Plan_A --> Selector
    Plan_B --> Selector
    Plan_C --> Selector
    Selector --> ExecPlan[最终执行计划]
    ExecPlan --> Executor[执行引擎]
    Executor --> Feedback[执行反馈数据]
    Feedback --> ML_Model
    Feedback --> Stats_Update[统计信息更新]

    style ML_Model fill:#e1f5fe
    style Cost_Pred fill:#e1f5fe
    style Feedback fill:#fff3e0

上图展示了 AI 查询优化的核心架构。关键机制包含三个层次：

第一层：学习型代价模型。传统 CBO 使用固定的代价公式（如 CPU_Cost + IO_Cost），参数权重硬编码。学习型代价模型将查询计划的物理特征（扫描行数、关联基数、索引选择性等）作为输入特征，用历史执行的真实耗时作为标签，训练回归模型预测代价。通过基准测试对比，在 TPC-H 基准上，学习型代价模型的代价预测误差从传统模型的 40%-60% 降低至 8%-15%。

第二层：计划选择强化学习。将查询优化视为序列决策问题——每一步选择 Join 顺序、访问路径、聚合策略都是一个动作。通过 Proximal Policy Optimization（PPO）等策略梯度方法，让智能体在历史查询负载上学习最优策略。Neo 数据库的研究表明，强化学习优化器在复杂多表关联场景下，计划选择质量可接近甚至超越人工调优。

第三层：执行反馈闭环。这是最关键的一环——每次查询执行后，将实际执行耗时、资源消耗、中间结果集大小等数据回传至模型，形成持续学习的闭环。这解决了传统优化器"一次决策，无法纠错"的根本缺陷。

三、生产级实现：基于 PostgreSQL 的学习型代价估算模块

以下代码展示了一个可嵌入 PostgreSQL 优化器钩子的学习型代价估算模块，使用轻量级梯度提升模型替代传统代价公式：

import numpy as np
import json
import logging
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field
import lightgbm as lgb
from pathlib import Path

logger = logging.getLogger("ai_query_optimizer")


@dataclass
class PlanFeatures:
    """查询计划的物理特征向量"""
    scan_rows: int = 0              # 预估扫描行数
    join_cardinality: int = 0       # 关联基数估算
    index_selectivity: float = 0.0  # 索引选择性 [0, 1]
    filter_selectivity: float = 0.0 # 过滤选择性 [0, 1]
    num_joins: int = 0              # 关联表数量
    num_group_keys: int = 0         # 分组键数量
    has_sort: bool = False          # 是否包含排序
    has_hash_agg: bool = False      # 是否使用哈希聚合
    has_merge_join: bool = False    # 是否使用归并关联
    table_size_mb: float = 0.0      # 涉及表的总大小（MB）
    index_hit_ratio: float = 0.0    # 索引命中率

    def to_vector(self) -> List[float]:
        """将特征转为模型输入向量"""
        return [
            np.log1p(self.scan_rows),        # 对数变换缓解长尾分布
            np.log1p(self.join_cardinality),
            self.index_selectivity,
            self.filter_selectivity,
            self.num_joins,
            self.num_group_keys,
            float(self.has_sort),
            float(self.has_hash_agg),
            float(self.has_merge_join),
            np.log1p(self.table_size_mb),
            self.index_hit_ratio,
        ]


@dataclass
class ExecutionFeedback:
    """执行反馈数据"""
    plan_features: PlanFeatures
    actual_latency_ms: float         # 实际执行耗时（毫秒）
    actual_rows: int                 # 实际返回行数
    actual_io_reads: int             # 实际 IO 读取次数
    cpu_time_ms: float               # CPU 时间（毫秒）
    plan_hash: str = ""              # 计划指纹，用于去重


class LearnedCostModel:
    """
    学习型代价模型：用梯度提升树替代传统代价公式
    支持在线增量训练，持续从执行反馈中学习
    """

    def __init__(self, model_path: Optional[str] = None):
        self.model: Optional[lgb.LGBMRegressor] = None
        self.feature_buffer: List[Tuple[List[float], float]] = []
        self.buffer_size = 500          # 累积 500 条反馈后触发增量训练
        self.min_train_samples = 200    # 最少训练样本数
        self.model_path = model_path or "/tmp/learned_cost_model.txt"
        self._load_model()

    def _load_model(self) -> None:
        """加载已保存的模型"""
        model_file = Path(self.model_path)
        if model_file.exists():
            try:
                self.model = lgb.Booster(model_file=str(model_file))
                logger.info("已加载学习型代价模型: %s", self.model_path)
            except Exception as e:
                logger.warning("模型加载失败，将使用冷启动模式: %s", e)
                self.model = None

    def predict_cost(self, features: PlanFeatures) -> float:
        """
        预测查询计划的执行代价（毫秒）
        若模型未就绪，回退到传统代价公式
        """
        if self.model is None:
            # 回退策略：基于特征的经验公式
            fallback = (
                np.log1p(features.scan_rows) * 0.01
                + np.log1p(features.join_cardinality) * 0.05
                + features.table_size_mb * 0.1
            )
            logger.debug("模型未就绪，回退代价: %.2f ms", fallback)
            return max(fallback, 0.1)

        vector = np.array([features.to_vector()])
        predicted = self.model.predict(vector)[0]
        # 代价预测值必须为正数，防止优化器选择异常
        return max(float(predicted), 0.1)

    def record_feedback(self, feedback: ExecutionFeedback) -> None:
        """记录执行反馈，累积后触发增量训练"""
        self.feature_buffer.append(
            (feedback.plan_features.to_vector(), feedback.actual_latency_ms)
        )
        logger.debug(
            "记录反馈: 实际耗时=%.2fms, 缓冲区=%d/%d",
            feedback.actual_latency_ms,
            len(self.feature_buffer),
            self.buffer_size,
        )

        if len(self.feature_buffer) >= self.buffer_size:
            self._incremental_train()

    def _incremental_train(self) -> None:
        """增量训练：基于累积的执行反馈更新模型"""
        if len(self.feature_buffer) < self.min_train_samples:
            logger.info("训练样本不足 (%d < %d)，跳过增量训练",
                        len(self.feature_buffer), self.min_train_samples)
            return

        X = np.array([f for f, _ in self.feature_buffer])
        y = np.array([l for _, l in self.feature_buffer])

        # 对延迟取对数，使目标分布更接近正态
        y_log = np.log1p(y)

        train_params = {
            "objective": "regression",
            "metric": "mae",
            "learning_rate": 0.05,
            "num_leaves": 31,
            "min_child_samples": 20,
            "feature_fraction": 0.8,
            "verbose": -1,
        }

        try:
            new_model = lgb.train(
                train_params,
                lgb.Dataset(X, y_log),
                num_boost_round=100,
                # 若已有模型，以已有模型为基础继续训练
                init_model=self.model,
            )
            self.model = new_model
            self.model.save_model(self.model_path)
            # 训练完成后清空缓冲区，保留最近 100 条用于下次训练的预热
            self.feature_buffer = self.feature_buffer[-100:]
            logger.info("增量训练完成，模型已保存至 %s", self.model_path)
        except Exception as e:
            logger.error("增量训练失败: %s", e)
            # 训练失败不影响在线服务，保留缓冲区下次重试


class QueryPlanAdvisor:
    """
    查询计划顾问：对比传统代价与学习型代价，
    当两者偏差超过阈值时发出优化建议
    """

    def __init__(self, cost_model: LearnedCostModel):
        self.cost_model = cost_model
        self.deviation_threshold = 2.0  # 偏差倍数阈值

    def compare_plans(
        self,
        plans: Dict[str, PlanFeatures],
        traditional_costs: Dict[str, float],
    ) -> Optional[str]:
        """
        对比多个候选计划的传统代价与学习型代价
        返回学习型模型推荐的计划 ID，若偏差不显著则返回 None
        """
        learned_costs = {}
        for plan_id, features in plans.items():
            learned_costs[plan_id] = self.cost_model.predict_cost(features)

        # 找到传统模型和学习模型各自的最优计划
        trad_best = min(traditional_costs, key=traditional_costs.get)
        learned_best = min(learned_costs, key=learned_costs.get)

        if trad_best == learned_best:
            logger.info("传统模型与学习模型选择一致: %s", trad_best)
            return None

        # 计算代价偏差倍数
        trad_cost_of_learned_best = traditional_costs[learned_best]
        trad_cost_of_trad_best = traditional_costs[trad_best]
        deviation = trad_cost_of_learned_best / max(trad_cost_of_trad_best, 0.001)

        if deviation > self.deviation_threshold:
            logger.warning(
                "学习模型推荐 %s，传统模型推荐 %s，偏差 %.1fx，建议人工审核",
                learned_best, trad_best, deviation,
            )
            return None  # 偏差过大，保守不切换

        logger.info(
            "学习模型推荐 %s (代价 %.2fms)，传统模型推荐 %s (代价 %.2fms)",
            learned_best, learned_costs[learned_best],
            trad_best, traditional_costs[trad_best],
        )
        return learned_best


# ============ 使用示例 ============

def demo():
    """演示学习型代价模型的完整工作流"""
    cost_model = LearnedCostModel()
    advisor = QueryPlanAdvisor(cost_model)

    # 模拟三个候选查询计划
    plans = {
        "plan_nl": PlanFeatures(
            scan_rows=10_000_000, join_cardinality=5_000_000,
            index_selectivity=0.01, num_joins=3,
            table_size_mb=500.0, index_hit_ratio=0.3,
        ),
        "plan_hash": PlanFeatures(
            scan_rows=10_000_000, join_cardinality=2_000_000,
            index_selectivity=0.01, num_joins=3,
            has_hash_agg=True, table_size_mb=500.0, index_hit_ratio=0.3,
        ),
        "plan_merge": PlanFeatures(
            scan_rows=10_000_000, join_cardinality=2_000_000,
            index_selectivity=0.01, num_joins=3,
            has_merge_join=True, has_sort=True,
            table_size_mb=500.0, index_hit_ratio=0.8,
        ),
    }

    # 传统代价估算（模拟 PostgreSQL CBO 输出）
    traditional_costs = {
        "plan_nl": 150.0,
        "plan_hash": 200.0,
        "plan_merge": 180.0,
    }

    recommended = advisor.compare_plans(plans, traditional_costs)
    print(f"推荐计划: {recommended or '维持传统模型选择'}")

    # 模拟执行反馈：Hash Join 实际执行最快
    feedback = ExecutionFeedback(
        plan_features=plans["plan_hash"],
        actual_latency_ms=85.0,
        actual_rows=1_800_000,
        actual_io_reads=12_000,
        cpu_time_ms=62.0,
        plan_hash="a1b2c3d4",
    )
    cost_model.record_feedback(feedback)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    demo()

关键设计决策说明：代价预测使用 log1p 变换处理延迟的长尾分布，这是因为在生产环境中查询耗时跨越毫秒到分钟量级，直接回归会导致模型偏向长尾样本。增量训练采用 LightGBM 的 init_model 机制，避免全量重训练的开销。偏差阈值设为 2.0 倍——当学习模型与传统模型推荐不一致且代价偏差超过 2 倍时，保守地不切换计划，防止模型误判引发性能回退。

四、学习型优化器的信任边界与工程权衡

AI 查询优化并非银弹，在工程落地中存在明确的边界条件与权衡：

模型冷启动问题。学习型代价模型在训练数据不足时，预测精度远低于传统公式。上述实现中通过回退策略缓解，但回退本身意味着系统在冷启动阶段无法提供优化价值。在生产环境中，至少需要积累 200-500 条带有真实执行耗时的查询样本，模型才能超越传统代价公式。这意味着新业务上线的前 1-2 周内，AI 优化器基本处于旁观状态。

推理延迟的额外开销。传统代价估算在微秒级完成，而模型推理即使使用 LightGBM 也需要 0.1-1ms。对于 OLTP 场景中大量短查询，这个额外开销占比可能达到 10%-30%，得不偿失。因此 AI 查询优化更适合 OLAP 或混合负载中的复杂查询——执行时间在百毫秒以上的查询，1ms 的推理开销可以忽略不计。

数据分布漂移。模型基于历史数据训练，当业务模式发生剧变（如大促、数据归档、表结构变更）时，模型预测精度会急剧下降。需要建立分布漂移检测机制——当连续 N 条查询的实际耗时与预测耗时的偏差超过阈值时，自动触发模型重训练。

可解释性缺失。传统代价公式可以逐项拆解为 CPU 代价、IO 代价、网络代价，便于 DBA 理解和调优。而学习型模型是黑盒，当它做出与传统模型相悖的推荐时，DBA 无法判断其合理性。在生产环境中，建议采用"影子模式"——AI 优化器仅做推荐不实际切换，积累置信度后再逐步放开。

训练数据污染风险。如果执行反馈中混入了因锁等待、网络抖动导致的异常耗时数据，模型会学习到错误的代价映射。必须在反馈采集层增加异常值过滤——剔除耗时偏离中位数 3 倍标准差以上的样本。

五、总结

AI 驱动的查询优化代表了数据库内核从"规则驱动"向"数据驱动"的进化方向。其核心价值在于：通过执行反馈闭环，让优化器从历史错误中学习，逐步逼近真实代价分布。但工程落地必须正视冷启动、推理开销、分布漂移和可解释性这四道门槛。务实的落地路线是：先在 OLAP 场景的复杂查询上以影子模式部署，积累足够的执行反馈并验证模型精度后，再逐步切换为主动推荐模式。对于 OLTP 短查询，传统代价公式仍是更可靠的选择。存储引擎的优化没有捷径，AI 是工具而非替代——最终决定查询性能的，仍然是对数据分布和访问模式的深刻理解。