AI 智能查询优化：从语义理解到执行计划自动改写

一、一条慢查询的 N 种写法，优化器只理解语法不理解意图

同一条业务查询，三种 SQL 写法，执行时间分别为 45 秒、3 秒、0.2 秒：

-- 写法1: 子查询嵌套, 优化器无法下推条件, 45秒
SELECT * FROM orders WHERE user_id IN (
    SELECT user_id FROM users WHERE region = '华东'
);

-- 写法2: JOIN 改写, 条件可下推, 3秒
SELECT o.* FROM orders o JOIN users u ON o.user_id = u.id
WHERE u.region = '华东';

-- 写法3: JOIN + 覆盖索引 + 分区裁剪, 0.2秒
SELECT o.order_id, o.amount FROM orders o
FORCE INDEX (idx_user_create)
JOIN users u FORCE INDEX (idx_region) ON o.user_id = u.id
WHERE u.region = '华东' AND o.create_time >= '2025-01-01';

三种写法的业务语义完全相同，但性能差距 225 倍。传统优化器只理解语法结构，不理解查询意图。AI 智能查询优化的目标：理解查询的语义意图，自动将低效写法改写为高效等价形式。

二、语义等价改写的底层机制

2.1 查询改写的等价规则

AI 查询优化的核心是一组语义等价改写规则，每条规则保证改写前后结果集不变：

规则	改写前	改写后	性能收益来源
子查询转 JOIN	IN (SELECT ...)	JOIN	条件下推, 减少中间结果
谓词下推	先 JOIN 后过滤	先过滤后 JOIN	减少 JOIN 行数
投影裁剪	SELECT *	SELECT 必要列	减少数据搬运, 覆盖索引
分区裁剪	全分区扫描	只扫描相关分区	减少扫描数据量
常量折叠	WHERE 1=1 AND ...	WHERE ...	减少无效计算
冗余 JOIN 消除	JOIN 无用表	移除 JOIN	减少连接操作
聚合下推	先 JOIN 后 GROUP BY	先 GROUP BY 后 JOIN	减少聚合数据量

2.2 AI 增强的改写决策

传统改写规则是确定性的：满足条件就改写。但某些改写是否有效取决于数据分布。例如，子查询转 JOIN 在子查询结果集小时有效，结果集大时可能更慢。

AI 增强的改写决策：用模型预测改写后的执行代价，只有预测代价更低时才执行改写。

flowchart TB
    A[原始 SQL] --> B[解析为 AST]
    B --> C[枚举可应用的改写规则]
    C --> D[生成候选改写 SQL]
    D --> E[代价预测模型]
    E --> F{预测代价 < 原始代价?}
    F -->|是| G[应用改写]
    F -->|否| H[保留原始]
    G --> I[输出优化后 SQL]
    H --> I

2.3 语义解析与意图识别

AI 查询优化需要理解 SQL 的语义意图，而非仅做语法匹配。关键步骤：

1. SQL 解析：将 SQL 文本解析为 AST（抽象语法树）
2. 语义标注：识别查询类型（点查/范围/聚合/分析）、表间关系、谓词类型
3. 意图分类：将查询归类到已知模式（如"按时间范围查某区域订单"）
4. 模板匹配：对已知模式应用预定义的最优写法模板

三、生产级 SQL 智能改写引擎

3.1 基于规则 + 模型混合的改写引擎

import re
from dataclasses import dataclass
from typing import List, Optional, Tuple
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class RewriteType(Enum):
    SUBQUERY_TO_JOIN = 'subquery_to_join'
    PREDICATE_PUSHDOWN = 'predicate_pushdown'
    PROJECTION_PRUNE = 'projection_prune'
    PARTITION_PRUNE = 'partition_prune'
    CONSTANT_FOLD = 'constant_fold'
    REDUNDANT_JOIN_ELIM = 'redundant_join_elim'
    AGGREGATE_PUSHDOWN = 'aggregate_pushdown'


@dataclass
class RewriteRule:
    """改写规则定义"""
    rule_type: RewriteType
    name: str
    description: str
    # 改写条件: 返回 True 时可应用
    condition: callable
    # 改写动作: 返回改写后的 SQL
    action: callable
    # 预估收益倍数
    estimated_benefit: float


@dataclass
class RewriteResult:
    """改写结果"""
    original_sql: str
    rewritten_sql: str
    applied_rules: List[str]
    estimated_improvement: float
    confidence: float  # 改写置信度


class SQLRewriteEngine:
    """SQL 智能改写引擎, 规则 + 模型混合"""

    def __init__(self):
        self.rules: List[RewriteRule] = []
        self._register_default_rules()

    def _register_default_rules(self):
        """注册默认改写规则"""
        # 规则1: IN 子查询转 JOIN
        self.rules.append(RewriteRule(
            rule_type=RewriteType.SUBQUERY_TO_JOIN,
            name='in_subquery_to_join',
            description='将 IN (SELECT ...) 子查询改写为 JOIN',
            condition=self._is_in_subquery,
            action=self._rewrite_in_subquery_to_join,
            estimated_benefit=5.0,
        ))

        # 规则2: 投影裁剪
        self.rules.append(RewriteRule(
            rule_type=RewriteType.PROJECTION_PRUNE,
            name='select_star_prune',
            description='将 SELECT * 改写为只查询必要列',
            condition=self._is_select_star,
            action=self._rewrite_select_star,
            estimated_benefit=2.0,
        ))

        # 规则3: 常量折叠
        self.rules.append(RewriteRule(
            rule_type=RewriteType.CONSTANT_FOLD,
            name='constant_fold',
            description='消除 WHERE 1=1 等恒真条件',
            condition=self._has_constant_condition,
            action=self._rewrite_constant_fold,
            estimated_benefit=1.1,
        ))

        # 规则4: 冗余 DISTINCT 消除
        self.rules.append(RewriteRule(
            rule_type=RewriteType.REDUNDANT_JOIN_ELIM,
            name='redundant_distinct_elim',
            description='当 SELECT 列包含主键时, DISTINCT 无意义',
            condition=self._has_redundant_distinct,
            action=self._rewrite_distinct_elim,
            estimated_benefit=1.3,
        ))

    def _is_in_subquery(self, sql: str) -> bool:
        """检测是否包含 IN 子查询"""
        return bool(re.search(r'\bIN\s*\(\s*SELECT\b', sql, re.IGNORECASE))

    def _rewrite_in_subquery_to_join(self, sql: str) -> str:
        """将 IN (SELECT ...) 改写为 JOIN"""
        # 匹配模式: WHERE col IN (SELECT col2 FROM table2 WHERE ...)
        pattern = r'(\w+)\s+IN\s*\(\s*SELECT\s+(\w+)\s+FROM\s+(\w+)(?:\s+WHERE\s+(.+?))?\s*\)'
        match = re.search(pattern, sql, re.IGNORECASE)

        if not match:
            return sql

        outer_col = match.group(1)
        inner_col = match.group(2)
        inner_table = match.group(3)
        inner_where = match.group(4)

        # 构造 JOIN 改写
        # 找到外层表名
        from_match = re.search(r'FROM\s+(\w+)', sql, re.IGNORECASE)
        if not from_match:
            return sql
        outer_table = from_match.group(1)

        # 移除 IN 子查询条件
        new_sql = re.sub(
            r'\b' + outer_col + r'\s+IN\s*\(\s*SELECT\s+\w+\s+FROM\s+\w+(?:\s+WHERE\s+.+?)?\s*\)',
            f'1=1',  # 占位, 后续替换
            sql, flags=re.IGNORECASE
        )

        # 添加 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}'

        new_sql = re.sub(
            r'FROM\s+' + outer_table,
            f'FROM {outer_table} {join_clause}',
            new_sql, flags=re.IGNORECASE
        )

        # 清理占位条件
        new_sql = re.sub(r'WHERE\s+1=1\s+AND\s+', 'WHERE ', new_sql, flags=re.IGNORECASE)
        new_sql = re.sub(r'\s+AND\s+1=1', '', new_sql, flags=re.IGNORECASE)
        new_sql = re.sub(r'WHERE\s+1=1', '', new_sql, flags=re.IGNORECASE)

        return new_sql

    def _is_select_star(self, sql: str) -> bool:
        return bool(re.search(r'SELECT\s+\*', sql, re.IGNORECASE))

    def _rewrite_select_star(self, sql: str) -> str:
        """将 SELECT * 改写为 SELECT 必要列 (需配合元数据)"""
        # 简化实现: 提示用户指定列, 无法自动推断所有必要列
        # 生产环境需结合 information_schema 获取表结构
        logger.warning("SELECT * 改写需要指定必要列, 建议手动优化")
        return sql

    def _has_constant_condition(self, sql: str) -> bool:
        return bool(re.search(r'\b1\s*=\s*1\b', sql, re.IGNORECASE))

    def _rewrite_constant_fold(self, sql: str) -> str:
        """消除恒真条件"""
        new_sql = re.sub(r'\bWHERE\s+1\s*=\s*1\s+AND\s+', 'WHERE ', sql, flags=re.IGNORECASE)
        new_sql = re.sub(r'\s+AND\s+1\s*=\s*1\b', '', new_sql, flags=re.IGNORECASE)
        new_sql = re.sub(r'\bWHERE\s+1\s*=\s*1\b', '', new_sql, flags=re.IGNORECASE)
        return new_sql

    def _has_redundant_distinct(self, sql: str) -> bool:
        return bool(re.search(r'SELECT\s+DISTINCT\b', sql, re.IGNORECASE))

    def _rewrite_distinct_elim(self, sql: str) -> str:
        """当 SELECT 列包含主键时, DISTINCT 无意义"""
        # 简化: 移除 DISTINCT (生产环境需检查是否包含主键)
        return re.sub(r'SELECT\s+DISTINCT\b', 'SELECT', sql, flags=re.IGNORECASE)

    def rewrite(self, sql: str) -> RewriteResult:
        """执行全部可应用的改写规则"""
        current_sql = sql
        applied_rules = []
        total_benefit = 1.0

        for rule in self.rules:
            if rule.condition(current_sql):
                new_sql = rule.action(current_sql)
                if new_sql != current_sql:
                    current_sql = new_sql
                    applied_rules.append(rule.name)
                    total_benefit *= rule.estimated_benefit
                    logger.info(f"应用规则 {rule.name}: {rule.description}")

        return RewriteResult(
            original_sql=sql,
            rewritten_sql=current_sql,
            applied_rules=applied_rules,
            estimated_improvement=round(total_benefit, 2),
            confidence=min(0.95, 0.7 + 0.05 * len(applied_rules)),
        )


class QueryPatternClassifier:
    """查询模式分类器, 识别查询意图并匹配最优模板"""

    # 已知查询模式及其最优写法模板
    PATTERNS = {
        'time_range_query': {
            'description': '按时间范围查询',
            'keywords': ['BETWEEN', '>=', 'create_time', 'date'],
            'optimization': '添加分区裁剪 + 时间索引 + 投影裁剪',
        },
        'user_behavior_query': {
            'description': '用户行为分析',
            'keywords': ['GROUP BY', 'user_id', 'COUNT', 'HAVING'],
            'optimization': '预聚合表 + 覆盖索引',
        },
        'multi_table_join': {
            'description': '多表关联查询',
            'keywords': ['JOIN', 'LEFT JOIN', 'ON'],
            'optimization': '小表驱动大表 + 谓词下推 + 投影裁剪',
        },
    }

    def classify(self, sql: str) -> List[dict]:
        """分类查询模式"""
        sql_upper = sql.upper()
        matched = []

        for pattern_name, pattern_info in self.PATTERNS.items():
            keyword_hits = sum(1 for kw in pattern_info['keywords'] if kw.upper() in sql_upper)
            if keyword_hits >= len(pattern_info['keywords']) * 0.5:
                matched.append({
                    'pattern': pattern_name,
                    'description': pattern_info['description'],
                    'optimization': pattern_info['optimization'],
                    'match_score': keyword_hits / len(pattern_info['keywords']),
                })

        return sorted(matched, key=lambda x: x['match_score'], reverse=True)

3.2 改写效果验证框架

class RewriteValidator:
    """改写结果验证器, 确保语义等价"""

    @staticmethod
    def validate_equivalence(original_sql: str, rewritten_sql: str,
                              mysql_config: dict) -> dict:
        """验证改写前后结果集是否一致"""
        result = {'equivalent': False, 'original_count': 0, 'rewritten_count': 0}

        try:
            import pymysql
            with pymysql.connect(**mysql_config) as conn:
                with conn.cursor() as cur:
                    # 比较结果行数
                    cur.execute(f"SELECT COUNT(*) FROM ({original_sql}) t")
                    result['original_count'] = cur.fetchone()[0]

                    cur.execute(f"SELECT COUNT(*) FROM ({rewritten_sql}) t")
                    result['rewritten_count'] = cur.fetchone()[0]

                    result['equivalent'] = (result['original_count'] == result['rewritten_count'])

                    # 进一步: 比较排序后的结果集 MD5 (仅小结果集)
                    if result['equivalent'] and result['original_count'] <= 10000:
                        cur.execute(f"SELECT MD5(GROUP_CONCAT(hash)) FROM ("
                                    f"SELECT MD5(CONCAT_WS('|', t.*)) as hash FROM ({original_sql}) t "
                                    f"ORDER BY hash) t2")
                        orig_hash = cur.fetchone()[0]

                        cur.execute(f"SELECT MD5(GROUP_CONCAT(hash)) FROM ("
                                    f"SELECT MD5(CONCAT_WS('|', t.*)) as hash FROM ({rewritten_sql}) t "
                                    f"ORDER BY hash) t2")
                        rewrite_hash = cur.fetchone()[0]

                        result['equivalent'] = (orig_hash == rewrite_hash)

        except Exception as e:
            result['error'] = str(e)

        return result

四、AI 查询优化的边界与架构妥协

4.1 语义等价的不完备性

SQL 语义等价判定在一般情况下是不可判定的。NULL 语义、DISTINCT 与 GROUP BY 的细微差异、窗口函数的边界行为，都可能导致改写后结果不一致。改写引擎必须对每条改写规则做严格的等价性证明，而非依赖测试覆盖。

4.2 改写爆炸

N 条可应用规则，排列组合产生 2^N 种改写方案。穷举所有组合的代价指数级增长。生产方案：贪心策略——按预估收益排序，依次应用，不做回溯。可能错过全局最优，但计算复杂度可控。

4.3 数据分布依赖

子查询转 JOIN 是否有效，取决于子查询结果集大小。改写引擎在无统计信息时无法判断，只能保守地不应用。需要与统计信息系统集成，获取实时 ndv 和行数估算。

4.4 禁用场景

• DML 语句（INSERT/UPDATE/DELETE）：改写可能改变副作用语义
• 包含用户自定义函数的查询：UDF 可能有副作用，改写不安全
• 存储过程和触发器内的 SQL：上下文依赖，无法独立改写
• 强制 Hint 的查询：用户已显式指定执行计划，改写违反用户意图

五、总结

AI 智能查询优化的核心是语义等价改写：理解查询意图，将低效写法自动转换为高效等价形式。改写规则保证结果集不变，AI 模型预测改写收益决定是否应用。生产落地的关键挑战是语义等价的严格保证——任何改写规则都必须有数学证明，而非仅靠测试验证。务实的路径是：先实现确定性规则改写（子查询转 JOIN、投影裁剪、常量折叠），再用 AI 模型做改写决策（何时应用、应用顺序），最后用结果集对比框架做上线前的等价性校验。AI 查询优化不是替代 DBA，而是将 DBA 的经验规则自动化，让每条 SQL 都能享受专家级优化。