大模型辅助的 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 重写仅作为规则引擎的补充。建议从高频慢查询开始优化，逐步积累重写规则库，并建立重写效果的量化追踪机制。