大模型辅助的 SQL 重写优化:从执行计划分析到语义等价变换的工程方案
大模型辅助的 SQL 重写优化:从执行计划分析到语义等价变换的工程方案
一、SQL 优化的经验瓶颈:为什么"改写 SQL"是 DBA 的手艺活
同一条查询逻辑,不同的 SQL 写法可能产生数十倍的性能差异。一个经典的例子:SELECT * FROM orders WHERE YEAR(create_time) = 2026无法使用create_time上的索引,而SELECT * FROM orders WHERE create_time >= '2026-01-01' AND create_time < '2027-01-01'可以。这种"函数导致索引失效"的问题,DBA 凭经验可以识别,但更复杂的重写——如子查询转 JOIN、UNION 优化、窗口函数替代自连接——需要深厚的 SQL 功底和大量的试错。
更深层的问题是:SQL 重写需要保证语义等价——重写后的查询必须返回与原查询完全相同的结果。人工重写容易在边界条件下引入 Bug(如 NULL 值处理、重复行消除),而验证语义等价性本身就很困难。
二、SQL 重写优化架构:从执行计划分析到语义等价验证
大模型辅助的 SQL 重写核心思路是:分析执行计划中的性能瓶颈,基于规则和 LLM 生成语义等价的重写方案,再通过结果对比验证等价性。
flowchart TD A[原始 SQL] --> B[执行计划分析<br/>EXPLAIN] B --> C[瓶颈识别<br/>全表扫描/临时表/文件排序] C --> D[规则匹配<br/>已知重写模式] D --> E{规则命中?} E -->|命中| F[应用规则重写] E -->|未命中| G[LLM 语义重写] F --> H[语义等价验证<br/>结果集对比] G --> H H --> I{结果一致?} I -->|是| J[性能对比<br/>执行时间/资源消耗] I -->|否| K[丢弃重写方案] J --> L{性能提升?} L -->|是| M[推荐重写方案] L -->|否| K关键设计决策在于语义等价验证的可靠性。简单的结果集对比在数据量大时不可行,需要采样验证或形式化验证。
三、工程实现:规则引擎、LLM 重写与等价验证
3.1 规则引擎:已知重写模式
from dataclasses import dataclass from typing import List, Optional import re @dataclass class RewriteRule: name: str pattern: str # 正则匹配原始 SQL description: str rewrite_fn: callable # 重写函数 class SQLRewriteRuleEngine: rules: List[RewriteRule] = [] def register(self, rule: RewriteRule): self.rules.append(rule) def match(self, sql: str) -> List[RewriteRule]: matched = [] for rule in self.rules: if re.search(rule.pattern, sql, re.IGNORECASE): matched.append(rule) return matched # 规则1:函数导致索引失效 → 范围查询 def rewrite_function_on_indexed_column(sql: str) -> Optional[str]: """YEAR(col) = N → col >= N-01-01 AND col < (N+1)-01-01""" pattern = r'WHERE\s+YEAR\((\w+)\)\s*=\s*(\d{4})' match = re.search(pattern, sql, re.IGNORECASE) if not match: return None column = match.group(1) year = int(match.group(2)) replacement = ( f"WHERE {column} >= '{year}-01-01' " f"AND {column} < '{year + 1}-01-01'" ) return re.sub(pattern, replacement, sql, flags=re.IGNORECASE) # 规则2:子查询转 JOIN def rewrite_subquery_to_join(sql: str) -> Optional[str]: """WHERE col IN (SELECT ...) → JOIN""" pattern = (r'WHERE\s+(\w+)\.(\w+)\s+IN\s*' r'\(\s*SELECT\s+(\w+)\s+FROM\s+(\w+)(?:\s+WHERE\s+(.+?))?\s*\)') match = re.search(pattern, sql, re.IGNORECASE) if not match: return None outer_table = match.group(1) outer_col = match.group(2) inner_col = match.group(3) inner_table = match.group(4) inner_where = match.group(5) # 构造 JOIN 语句 join_clause = f"JOIN {inner_table} ON {outer_table}.{outer_col} = {inner_table}.{inner_col}" if inner_where: join_clause += f" AND {inner_where}" # 替换 WHERE IN 为 JOIN rewritten = re.sub( r'FROM\s+(\w+)\s+WHERE\s+' + re.escape(match.group(0).split('WHERE')[1]), f"FROM {outer_table} {join_clause} WHERE", sql, flags=re.IGNORECASE ) return rewritten # 规则3:OR 条件转 UNION ALL def rewrite_or_to_union(sql: str) -> Optional[str]: """WHERE a = 1 OR b = 2 → UNION ALL of two queries""" pattern = r'WHERE\s+(.+?)\s+OR\s+(.+?)(?:\s+ORDER|\s+LIMIT|\s*$)' match = re.search(pattern, sql, re.IGNORECASE) if not match: return None cond_a = match.group(1).strip() cond_b = match.group(2).strip() # 提取 FROM 之前的部分 select_part = sql[:sql.upper().index('FROM')].strip() from_part = sql[sql.upper().index('FROM'):sql.upper().index('WHERE')].strip() return f"{select_part} {from_part} WHERE {cond_a} UNION ALL {select_part} {from_part} WHERE {cond_b}"3.2 LLM 语义重写
class LLMSQLRewriter: def __init__(self, llm_client): self.llm = llm_client def rewrite(self, sql: str, explain_output: str) -> Optional[str]: prompt = f"""你是一个 SQL 优化专家。请重写以下 SQL 以提升性能,保持语义完全等价。 原始 SQL: {sql} 执行计划分析: {explain_output} 重写要求: 1. 保持语义等价(返回完全相同的结果集) 2. 消除全表扫描、临时表和文件排序 3. 优先使用 JOIN 替代子查询 4. 避免在索引列上使用函数 5. 只输出重写后的 SQL,不要解释 重写后的 SQL:""" response = self.llm.call(prompt, max_tokens=2000) # 提取 SQL 代码块 if '```sql' in response: return response.split('```sql')[1].split('```')[0].strip() if '```' in response: return response.split('```')[1].split('```')[0].strip() return response.strip()3.3 语义等价验证
class SemanticEquivalenceValidator: def __init__(self, db_connection): self.conn = db_connection def validate(self, original_sql: str, rewritten_sql: str, sample_size: int = 1000) -> bool: """通过采样对比验证语义等价""" # 方案1:小数据集全量对比 # 在测试环境中执行两条 SQL,对比结果集 # 方案2:采样对比 # 添加 LIMIT 后对比前 N 行 limited_original = f"{original_sql.rstrip(';')} LIMIT {sample_size}" limited_rewritten = f"{rewritten_sql.rstrip(';')} LIMIT {sample_size}" try: result_a = self.conn.execute(limited_original).fetchall() result_b = self.conn.execute(limited_rewritten).fetchall() # 排序后对比(顺序可能不同) sorted_a = sorted([tuple(r) for r in result_a]) sorted_b = sorted([tuple(r) for r in result_b]) return sorted_a == sorted_b except Exception as e: # 重写后的 SQL 执行失败,不等价 return False def validate_with_checksum(self, original_sql: str, rewritten_sql: str) -> bool: """通过校验和验证全量等价(适用于小表)""" checksum_a = self._compute_checksum(original_sql) checksum_b = self._compute_checksum(rewritten_sql) return checksum_a == checksum_b def _compute_checksum(self, sql: str) -> str: import hashlib cursor = self.conn.execute(sql) hasher = hashlib.md5() for row in cursor: hasher.update(str(tuple(row)).encode()) return hasher.hexdigest()四、SQL 重写的等价性风险与适用边界
NULL 值的语义差异:NOT IN (subquery)在子查询结果包含 NULL 时行为与NOT EXISTS不同——NOT IN返回空集,而NOT EXISTS返回正确结果。重写时必须考虑 NULL 语义,否则可能引入 Bug。
排序的隐式依赖:某些应用依赖查询结果的隐式排序(如 MySQL 在某些情况下按主键排序返回),但 SQL 标准不保证无ORDER BY时的排序。重写后排序可能变化,导致应用行为异常。
LLM 重写的不可靠性:LLM 可能生成语法正确但语义不等价的 SQL。例如,将LEFT JOIN重写为INNER JOIN会丢失不匹配的行。语义等价验证是必不可少的环节,但采样验证无法覆盖所有边界条件。
重写收益的上下文依赖:同一条重写规则在不同数据分布下效果不同。子查询转 JOIN 在小表驱动大表时性能提升明显,但在大表驱动小表时可能更差。重写推荐需要结合数据分布和执行计划做综合判断。
五、总结
大模型辅助 SQL 重写的本质是将"DBA 经验驱动的改写"转化为"规则匹配 + LLM 语义推导 + 等价验证"的系统化方案。本文方案的核心链路为:执行计划瓶颈识别 → 规则引擎匹配 → LLM 语义重写 → 采样等价验证 → 性能对比。落地时需重点关注三个原则:所有重写必须通过等价验证、优先使用规则引擎处理已知模式、LLM 重写仅作为规则引擎的补充。建议从高频慢查询开始优化,逐步积累重写规则库,并建立重写效果的量化追踪机制。
