AI 查询优化器：当机器学习遇上执行计划，数据库调优的新范式

一、数据库调优的"玄学"困境：为什么 DBA 靠经验吃饭

传统数据库查询优化依赖 DBA 的经验：看执行计划、调索引、改写 SQL、调整统计信息。这个过程高度依赖人的直觉和经验积累，且不同 DBA 可能给出截然不同的优化方案。更棘手的是，数据分布变化后，原本最优的执行计划可能变得次优——但 DBA 不可能 24 小时盯着每条 SQL 的执行计划变化。

AI 查询优化器的核心思路：用机器学习模型替代 DBA 的经验判断，自动预测最优执行计划。模型从历史查询的执行数据中学习"什么类型的查询在什么数据分布下用什么执行计划最快"，并在新查询到来时实时推荐最优方案。这不是取代 DBA，而是将 DBA 的经验编码为可扩展的算法。

二、AI 查询优化的技术架构：从代价模型到学习型优化器

flowchart TB
    A[SQL 查询输入] --> B[查询解析与特征提取]
    B --> C{优化路径}
    C -->|传统路径| D[基于代价的优化器 CBO]
    C -->|AI 路径| E[学习型优化器 LBO]

    D --> D1[统计信息估算]
    D1 --> D2[代价模型计算]
    D2 --> D3[枚举执行计划]
    D3 --> F[选择最低代价计划]

    E --> E1[查询编码: AST→向量]
    E1 --> E2[计划评分模型]
    E2 --> E3[Top-K 候选计划]
    E3 --> E4[快速验证: 实际执行采样]
    E4 --> F

    F --> G[执行查询]
    G --> H[收集执行统计]
    H --> E2

    style C fill:#ffd93d,color:#333
    style E fill:#4d96ff,color:#fff
    style H fill:#6bcb77,color:#fff

两种优化路径的对比：

• CBO（基于代价的优化器）：数据库内置的优化器，依赖统计信息（直方图、NDV、行数估算）计算每种执行计划的代价，选择代价最低的。问题在于统计信息可能过时，代价模型是简化的（无法精确预测缓存命中率、CPU 缓存行为等）。
• LBO（学习型优化器）：用模型预测执行计划的代价或直接推荐最优计划。模型输入是查询的特征编码（表名、谓词、连接条件），输出是计划的评分或排序。优势是可以学习到 CBO 无法建模的复杂模式（如特定数据分布下的最优 Join 顺序）。

三、学习型查询优化器实现

# 查询特征提取 — 将 SQL AST 编码为模型输入
import sqlparse
import numpy as np
from typing import List, Dict


class QueryEncoder:
    """将 SQL 查询编码为固定长度的特征向量"""

    def __init__(self, table_vocab: Dict[str, int],
                 column_vocab: Dict[str, int],
                 operator_vocab: Dict[str, int]):
        self.table_vocab = table_vocab
        self.column_vocab = column_vocab
        self.operator_vocab = operator_vocab

    def encode(self, sql: str) -> np.ndarray:
        """
        编码 SQL 查询为特征向量
        特征包括：表名、列名、操作符、谓词选择性、连接条件数
        """
        parsed = sqlparse.parse(sql)[0]

        features = []

        # 特征 1：涉及的表（multi-hot 编码）
        table_vec = np.zeros(len(self.table_vocab))
        for table, idx in self.table_vocab.items():
            if table.lower() in sql.lower():
                table_vec[idx] = 1
        features.extend(table_vec)

        # 特征 2：涉及的列（multi-hot 编码）
        column_vec = np.zeros(len(self.column_vocab))
        for col, idx in self.column_vocab.items():
            if col.lower() in sql.lower():
                column_vec[idx] = 1
        features.extend(column_vec)

        # 特征 3：操作符统计
        op_counts = np.zeros(len(self.operator_vocab))
        for op, idx in self.operator_vocab.items():
            op_counts[idx] = sql.count(op)
        features.extend(op_counts)

        # 特征 4：查询复杂度指标
        features.append(sql.count('JOIN'))       # 连接数
        features.append(sql.count('WHERE'))       # 过滤条件数
        features.append(sql.count('GROUP BY'))    # 聚合数
        features.append(sql.count('ORDER BY'))    # 排序数
        features.append(sql.count('SUBQUERY'))    # 子查询数

        return np.array(features, dtype=np.float32)


# 执行计划评分模型
import torch
import torch.nn as nn


class PlanScoringModel(nn.Module):
    """
    执行计划评分模型
    输入：查询特征 + 执行计划特征
    输出：预测执行时间（ms）
    """

    def __init__(self, query_dim: int, plan_dim: int,
                 hidden_dim: int = 256):
        super().__init__()

        # 查询编码器
        self.query_encoder = nn.Sequential(
            nn.Linear(query_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

        # 计划编码器
        self.plan_encoder = nn.Sequential(
            nn.Linear(plan_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

        # 联合评分头
        self.scorer = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, 1),  # 输出预测执行时间
        )

    def forward(self, query_features: torch.Tensor,
                plan_features: torch.Tensor) -> torch.Tensor:
        q_enc = self.query_encoder(query_features)
        p_enc = self.plan_encoder(plan_features)

        # 拼接查询和计划特征
        combined = torch.cat([q_enc, p_enc], dim=-1)
        predicted_time = self.scorer(combined)

        # 输出对数时间（执行时间跨度大，用对数空间更稳定）
        return predicted_time.squeeze(-1)


# 执行计划特征提取
class PlanEncoder:
    """将执行计划编码为特征向量"""

    def encode(self, plan: Dict) -> np.ndarray:
        """
        编码执行计划的关键特征
        包括：Join 类型、扫描方式、索引使用、预估行数
        """
        features = []

        # Join 类型统计
        features.append(plan.get('hash_join_count', 0))
        features.append(plan.get('nested_loop_count', 0))
        features.append(plan.get('merge_join_count', 0))

        # 扫描方式统计
        features.append(plan.get('index_scan_count', 0))
        features.append(plan.get('seq_scan_count', 0))
        features.append(plan.get('bitmap_scan_count', 0))

        # 预估代价
        features.append(np.log1p(plan.get('estimated_cost', 0)))
        features.append(np.log1p(plan.get('estimated_rows', 0)))

        # 排序和聚合
        features.append(plan.get('sort_count', 0))
        features.append(plan.get('aggregate_count', 0))

        return np.array(features, dtype=np.float32)

# AI 优化器服务 — 集成到数据库查询流程
class AIQueryOptimizer:
    """AI 查询优化器：在 CBO 基础上提供候选计划评分"""

    def __init__(self, scoring_model, query_encoder, plan_encoder):
        self.model = scoring_model
        self.query_encoder = query_encoder
        self.plan_encoder = plan_encoder

    def recommend_plan(self, sql: str,
                       candidate_plans: List[Dict]) -> Dict:
        """
        从候选执行计划中选择 AI 推荐的最优计划
        candidate_plans: CBO 生成的多个候选计划
        """
        query_feat = self.query_encoder.encode(sql)
        query_tensor = torch.tensor(query_feat).unsqueeze(0)

        best_plan = None
        best_score = float('inf')

        for plan in candidate_plans:
            plan_feat = self.plan_encoder.encode(plan)
            plan_tensor = torch.tensor(plan_feat).unsqueeze(0)

            with torch.no_grad():
                predicted_time = self.model(query_tensor, plan_tensor)

            if predicted_time.item() < best_score:
                best_score = predicted_time.item()
                best_plan = plan

        return {
            'recommended_plan': best_plan,
            'predicted_time_ms': np.exp(best_score),  # 还原对数时间
            'confidence': self._compute_confidence(candidate_plans),
        }

    @staticmethod
    def _compute_confidence(plans: List[Dict]) -> float:
        """计算推荐置信度：候选计划间差异越大，置信度越高"""
        if len(plans) < 2:
            return 0.5
        # 简化：基于代价差异的置信度
        costs = [p.get('estimated_cost', 0) for p in plans]
        cost_range = max(costs) - min(costs)
        return min(1.0, cost_range / (max(costs) + 1))

四、AI 查询优化的现实挑战

训练数据的冷启动：模型需要大量"查询 + 执行计划 + 实际执行时间"的三元组数据。新数据库上线时没有历史数据，模型无法训练。解决方案是用 CBO 的代价估算作为初始标签，随着真实执行数据的积累逐步替换。

数据分布漂移：模型在训练时的数据分布与线上可能不同（数据量增长、索引变更、业务模式变化）。模型需要持续在线学习，但在线学习的稳定性难以保证——新数据可能引入噪声，导致模型性能震荡。

安全性与可解释性：DBA 不信任黑箱模型的推荐——如果 AI 选择了全表扫描而非索引扫描，DBA 需要知道原因。模型必须提供可解释的推荐理由（如"该索引的选择性在当前数据分布下低于 0.01，全表扫描更优"），否则不会被采纳。

优化空间的限制：AI 优化器只能在 CBO 生成的候选计划中选择，无法创造 CBO 枚举不到的执行计划。真正的突破需要 AI 参与计划生成（如学习新的 Join 算法组合），但这需要深度修改数据库内核，工程难度极高。

五、总结

AI 查询优化器的核心原则：辅助而非替代 CBO、数据驱动决策、可解释推荐。落地路径：

1. 起步阶段：收集查询执行日志（SQL + 执行计划 + 实际耗时），建立训练数据集，训练计划评分模型。
2. 进阶阶段：在测试环境中部署 AI 优化器，对比 AI 推荐计划与 CBO 选择计划的实际执行时间，验证优化效果。
3. 成熟阶段：在生产环境中以"影子模式"运行 AI 优化器（推荐但不强制执行），积累信任后逐步切换。
4. 持续优化：建立在线学习管线，用最新执行数据持续更新模型，适应数据分布变化。

AI 查询优化不是一蹴而就的，而是从"辅助决策"到"自动优化"的渐进过程。先用数据证明 AI 的推荐比 CBO 更准，再谈替代。