基于强化学习的查询优化:让数据库自己学会选择执行路径

发布时间:2026/6/27 2:52:08
基于强化学习的查询优化:让数据库自己学会选择执行路径 基于强化学习的查询优化让数据库自己学会选择执行路径一、从规则到学习查询优化的范式迁移传统查询优化器依赖两个核心组件代价模型和搜索策略。代价模型估算每个执行计划的 IO 和 CPU 开销搜索策略在等价计划空间中寻找代价最低的方案。但这两个组件都有固有局限代价模型的精度受统计信息采样率限制搜索策略受启发式剪枝约束容易陷入局部最优。强化学习Reinforcement Learning, RL提供了一种全新的范式不显式建模代价函数而是让 Agent 通过与数据库环境交互学习哪种执行计划在哪种查询特征下表现最好。Agent 的奖励是查询的实际执行延迟——这是最直接的优化目标绕过了代价模型的精度瓶颈。但 RL 查询优化面临三个核心挑战状态空间巨大查询特征 数据统计的组合、奖励稀疏只有查询执行完才能获得延迟反馈、探索风险高劣质执行计划可能导致 OOM 或长时间阻塞。二、RL 查询优化的建模与架构flowchart TD subgraph 环境定义 A[状态 Statebr/查询特征向量br/ 数据统计信息] B[动作 Actionbr/JOIN 顺序选择br/ 索引选择br/ 算子选择] C[奖励 Rewardbr/-log(执行延迟/ms)br/OOM/超时 大负值] end subgraph RL Agent D[策略网络 Policy Networkbr/输入: 状态向量br/输出: 动作概率分布] E[价值网络 Value Networkbr/输入: 状态向量br/输出: 状态价值估计] end subgraph 训练循环 F[采样查询] G[Agent 选择执行计划] H[执行查询获取延迟] I[计算奖励] J[更新策略网络和价值网络] end A -- D A -- E D -- B B -- G G -- H H -- I I -- J J -- D J -- E C -.- I style D fill:#e1f5fe style E fill:#fff3e0 style C fill:#ffebee2.1 状态空间设计状态向量需要编码查询的结构特征和数据统计信息同时维度不能过高高维状态导致训练收敛慢特征类别具体特征维度查询结构表数量、JOIN 数量、谓词数量、聚合列数4数据统计各表行数log、各索引 Cardinality2N查询复杂度子查询深度、UNION 数量、排序列数3历史信息最近 5 次同类查询的平均延迟52.2 动作空间设计查询优化的动作空间是组合爆炸的。对于 N 表 JOIN仅 JOIN 顺序就有 N! 种可能。RL 方案采用分层决策降低动作空间第一层选择第一对 JOIN 的表N*(N-1)/2 种选择。第二层选择 JOIN 算法Hash Join / Nested Loop / Merge Join。第三层选择后续 JOIN 的表和算法。每层的动作空间控制在 100 以内避免策略网络的输出维度过高。2.3 奖励函数设计import math from typing import Optional def compute_reward( execution_time_ms: float, baseline_time_ms: float, is_oom: bool False, is_timeout: bool False, timeout_threshold_ms: float 60000.0 ) - float: 计算 RL Agent 的奖励值。 为什么用 -log 而非负延迟 延迟的绝对值差异大1ms vs 10000ms 直接用负延迟会导致策略偏向避免极端延迟 而忽略对中等延迟查询的优化。 -log 使奖励在延迟的对数尺度上均匀分布 让 Agent 对各延迟区间的优化同等重视。 if is_oom: # OOM 是最严重的惩罚查询失败且影响系统稳定性 return -100.0 if is_timeout: return -50.0 if execution_time_ms 0: execution_time_ms 0.1 # 防止 log(0) # 基础奖励延迟越低奖励越高 reward -math.log10(execution_time_ms) # 相对基线的改进奖励 # 为什么需要基线比较绝对延迟受数据量影响大 # 与基线比较能消除数据量的影响让 Agent 关注优化效果 if baseline_time_ms 0: improvement_ratio baseline_time_ms / execution_time_ms reward math.log2(improvement_ratio) return reward三、PPO 策略训练与安全探索3.1 基于 PPO 的策略优化import torch import torch.nn as nn import torch.optim as optim import numpy as np from typing import List, Tuple class QueryOptimizationPolicy(nn.Module): 查询优化策略网络输入查询状态输出动作概率。 为什么用 PPO 而非 DQN 1. 动作空间是连续的JOIN 顺序 算子选择的组合 DQN 的离散动作空间不适用 2. PPO 的策略梯度方法更稳定适合在线学习 3. PPO 的 clip 机制防止策略更新过大保证训练稳定性 def __init__( self, state_dim: int, action_dim: int, hidden_dim: int 256 ): super().__init__() self.shared_net nn.Sequential( nn.Linear(state_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), ) # 策略头输出动作概率 self.policy_head nn.Sequential( nn.Linear(hidden_dim, action_dim), nn.Softmax(dim-1) ) # 价值头输出状态价值 self.value_head nn.Linear(hidden_dim, 1) def forward( self, state: torch.Tensor ) - Tuple[torch.Tensor, torch.Tensor]: shared self.shared_net(state) action_probs self.policy_head(shared) state_value self.value_head(shared) return action_probs, state_value class PPOTrainer: PPO 训练器收集轨迹数据并更新策略网络。 def __init__( self, policy: QueryOptimizationPolicy, lr: float 3e-4, gamma: float 0.99, clip_epsilon: float 0.2, # GAE 参数平衡偏差和方差 gae_lambda: float 0.95 ): self.policy policy self.optimizer optim.Adam(policy.parameters(), lrlr) self.gamma gamma self.clip_epsilon clip_epsilon self.gae_lambda gae_lambda def compute_advantages( self, rewards: List[float], values: List[float], dones: List[bool] ) - List[float]: 计算广义优势估计GAE。 为什么用 GAE 而非简单的 TD 误差 GAE 通过 lambda 参数在偏差和方差之间插值 lambda0 时退化为 TD(0)低方差高偏差 lambda1 时退化为 Monte Carlo高方差低偏差 lambda0.95 是经验上的良好折中。 advantages [] gae 0 next_value 0 for t in reversed(range(len(rewards))): if dones[t]: next_value 0 gae 0 delta ( rewards[t] self.gamma * next_value * (1 - dones[t]) - values[t] ) gae delta self.gamma * self.gae_lambda * (1 - dones[t]) * gae advantages.insert(0, gae) next_value values[t] return advantages def update( self, states: torch.Tensor, actions: torch.Tensor, old_log_probs: torch.Tensor, returns: torch.Tensor, advantages: torch.Tensor ) - Dict: PPO 策略更新 # 归一化优势值稳定训练 advantages (advantages - advantages.mean()) / ( advantages.std() 1e-8 ) action_probs, values self.policy(states) dist torch.distributions.Categorical(action_probs) new_log_probs dist.log_prob(actions) # PPO clip 目标 ratio torch.exp(new_log_probs - old_log_probs) surr1 ratio * advantages surr2 torch.clamp( ratio, 1 - self.clip_epsilon, 1 self.clip_epsilon ) * advantages # 策略损失取 clip 后的较小值 policy_loss -torch.min(surr1, surr2).mean() # 价值损失MSE value_loss nn.MSELoss()(values.squeeze(), returns) # 熵正则化鼓励探索防止策略过早收敛 entropy dist.entropy().mean() entropy_coeff 0.01 total_loss ( policy_loss 0.5 * value_loss - entropy_coeff * entropy ) self.optimizer.zero_grad() total_loss.backward() # 梯度裁剪防止梯度爆炸 nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5) self.optimizer.step() return { policy_loss: policy_loss.item(), value_loss: value_loss.item(), entropy: entropy.item(), approx_kl: ( (new_log_probs - old_log_probs).mean().item() ) }3.2 安全探索机制class SafeExplorer: 安全探索器限制 RL Agent 的探索范围 防止选择导致 OOM 或超时的执行计划。 为什么需要安全探索RL 的探索可能选择极端劣质计划 在生产环境中直接执行会导致系统不可用。 def __init__( self, optimizer, # 传统优化器作为安全基线 max_latency_ratio: float 5.0, # 允许的最大延迟倍数 max_memory_mb: float 4096 # 单查询内存上限 ): self.optimizer optimizer self.max_latency_ratio max_latency_ratio self.max_memory_mb max_memory_mb def safe_select( self, query: str, rl_plan: ExecutionPlan, state: np.ndarray ) - ExecutionPlan: 安全选择在 RL 计划和传统优化器计划之间选择。 策略如果 RL 计划的预估代价不超过传统计划的 N 倍 则使用 RL 计划否则降级到传统计划。 # 获取传统优化器的计划作为安全基线 baseline_plan self.optimizer.optimize(query) # 估算 RL 计划的代价 rl_cost self._estimate_cost(rl_plan, state) baseline_cost self._estimate_cost(baseline_plan, state) if rl_cost baseline_cost * self.max_latency_ratio: logging.warning( fRL 计划代价 {rl_cost:.2f} 超过基线 f{baseline_cost:.2f} 的 {self.max_latency_ratio}x f降级到传统计划 ) return baseline_plan # 检查内存预估 estimated_memory self._estimate_memory(rl_plan, state) if estimated_memory self.max_memory_mb: logging.warning( fRL 计划预估内存 {estimated_memory:.0f}MB f超过上限 {self.max_memory_mb}MB降级 ) return baseline_plan return rl_plan def _estimate_cost( self, plan: ExecutionPlan, state: np.ndarray ) - float: 使用传统代价模型估算计划代价 # 委托给传统优化器的代价估算模块 return self.optimizer.estimate_cost(plan) def _estimate_memory( self, plan: ExecutionPlan, state: np.ndarray ) - float: 估算计划的内存消耗 return self.optimizer.estimate_memory(plan)四、RL 查询优化的落地瓶颈与适用边界4.1 训练数据的需求PPO 需要大量查询-执行计划-延迟三元组来训练策略网络。在生产环境中每条查询的执行延迟从毫秒到分钟不等收集足够的训练数据可能需要数周。离线训练使用历史查询日志可以加速冷启动但历史数据中的执行计划是由传统优化器选择的缺乏多样性。4.2 泛化性问题策略网络在训练数据覆盖的查询模式上表现良好但对未见过的查询模式如新增的 JOIN 组合可能做出错误决策。解决方案是混合策略对已知查询模式使用 RL 计划对未知模式降级到传统优化器。4.3 在线学习的风险RL 的在线学习需要持续执行查询并获取反馈。在生产环境中这意味着部分查询会被用作训练样本执行 RL 选择的计划。如果 RL 计划显著劣于传统计划这些查询的延迟会恶化。安全探索机制可以缓解但不能完全消除这一风险。4.4 适用边界场景RL 优化适用性推荐替代固定查询模式BI 报表极佳物化视图多表 JOIN6 表良好代价模型增强OLTP 点查不适用传统优化器查询模式频繁变化差泛化不足传统优化器离线分析场景良好可容忍训练开销手动调优五、总结基于强化学习的查询优化将执行计划选择建模为序贯决策问题通过与环境交互学习最优策略。PPO 算法的 clip 机制和 GAE 优势估计提供了稳定的训练框架安全探索机制防止了劣质计划对生产系统的影响。但 RL 查询优化的落地仍面临训练数据需求大、泛化性不足和在线学习风险三个核心挑战。当前阶段RL 优化更适合作为传统优化器的增强层而非替代品。落地路线建议第一步在测试环境中用历史查询日志进行离线训练验证策略网络的收敛性第二步以影子模式部署RL 选择计划但不执行仅记录与传统优化器的差异第三步对 RL 计划优于传统计划的查询模式逐步启用 RL 选择同时保持传统优化器作为降级兜底。