向量化引擎的AI辅助查询优化：从执行计划到向量化路径的智能推导

一、向量化执行的性能鸿沟：为什么"能跑"和"跑得快"之间差了一个编译器

向量化执行引擎（如 ClickHouse 的执行引擎）通过批量处理列式数据，将 CPU 缓存利用率提升数倍。但并非所有查询都能自动受益于向量化——复杂嵌套查询、多表 JOIN、UDF 调用等场景下，执行引擎可能退回到逐行处理模式，性能骤降 10-50 倍。问题在于：向量化路径的选择依赖查询的物理执行计划，而传统优化器基于规则（RBO）或简单代价模型（CBO）生成计划时，无法感知向量化执行的 CPU 缓存行为和 SIMD 指令利用率。

AI 辅助查询优化的核心价值在于：通过学习历史查询的执行特征，预测不同物理计划在向量化引擎上的真实性能，从而选择最优的向量化执行路径。这不再是"估算 I/O 成本"，而是"预测 CPU 行为"。

二、AI驱动的向量化路径选择架构

flowchart TB
    A[SQL 查询输入] --> B[逻辑计划解析]
    B --> C[候选物理计划生成]
    C --> D[特征提取]
    D --> E[AI 代价模型]
    E --> F[最优计划选择]
    F --> G[向量化执行引擎]

    subgraph 特征提取
        D1[算子类型分布] --> D
        D2[数据基数估计] --> D
        D3[内存访问模式] --> D
        D4[JOIN 顺序特征] --> D
    end

    subgraph AI 代价模型
        E1[执行时间预测] --> E
        E2[缓存命中率估计] --> E
        E3[SIMD 利用率评估] --> E
    end

    G --> H[执行反馈]
    H -->|在线学习| E

AI 代价模型接收候选物理计划的特征向量，输出预测的执行时间、缓存命中率和 SIMD 利用率。与传统 CBO 的区别在于：传统模型基于统计信息估算 I/O 和 CPU 成本，AI 模型基于历史执行数据学习"计划特征→真实性能"的映射关系。执行反馈回路使模型持续优化——每次查询执行后，真实性能数据作为训练样本更新模型。

三、AI代价模型与向量化路径选择的工程实现

3.1 查询计划特征提取器

import numpy as np
from dataclasses import dataclass
from typing import List

@dataclass
class PlanFeature:
    """物理执行计划的特征向量"""
    # 算子特征
    scan_count: int          # 扫描算子数量
    join_count: int          # JOIN 算子数量
    aggregate_count: int     # 聚合算子数量
    sort_count: int          # 排序算子数量
    filter_count: int        # 过滤算子数量

    # 数据特征
    total_input_rows: float  # 输入行数估计（log scale）
    max_cardinality: float   # 最大基数估计
    join_selectivity: float  # JOIN 选择率
    filter_selectivity: float # 过滤选择率

    # 向量化特征
    vectorizable_ops: int    # 可向量化算子数
    row_level_ops: int       # 逐行处理算子数
    simd_friendly: float     # SIMD 友好度（0-1）

    # 内存特征
    estimated_memory_mb: float
    spill_risk: float        # 溢出风险（0-1）

class PlanFeatureExtractor:
    """从物理执行计划中提取特征向量"""

    # 可向量化算子白名单
    VECTORIZABLE_OPS = {
        'TableScan', 'Filter', 'Projection', 'Aggregation',
        'HashJoin', 'MergeJoin', 'Limit'
    }

    # SIMD 友好的数据类型
    SIMD_TYPES = {'Int8', 'Int16', 'Int32', 'Int64', 'Float32', 'Float64'}

    def extract(self, physical_plan: dict) -> PlanFeature:
        """递归遍历计划树，提取特征"""
        ops = self._collect_operators(physical_plan)

        scan_ops = [o for o in ops if o['type'] == 'TableScan']
        join_ops = [o for o in ops if 'Join' in o['type']]
        agg_ops = [o for o in ops if o['type'] == 'Aggregation']
        sort_ops = [o for o in ops if o['type'] == 'Sort']
        filter_ops = [o for o in ops if o['type'] == 'Filter']

        vectorizable = sum(1 for o in ops if o['type'] in self.VECTORIZABLE_OPS)
        row_level = len(ops) - vectorizable

        # 估算 SIMD 友好度：可向量化算子中 SIMD 类型列的占比
        simd_columns = 0
        total_columns = 0
        for op in scan_ops:
            for col in op.get('columns', []):
                total_columns += 1
                if any(t in col.get('type', '') for t in self.SIMD_TYPES):
                    simd_columns += 1

        simd_friendly = simd_columns / max(total_columns, 1)

        return PlanFeature(
            scan_count=len(scan_ops),
            join_count=len(join_ops),
            aggregate_count=len(agg_ops),
            sort_count=len(sort_ops),
            filter_count=len(filter_ops),
            total_input_rows=np.log1p(
                sum(o.get('estimated_rows', 0) for o in scan_ops)
            ),
            max_cardinality=np.log1p(
                max((o.get('estimated_rows', 1) for o in ops), default=1)
            ),
            join_selectivity=np.mean([
                o.get('selectivity', 0.5) for o in join_ops
            ]) if join_ops else 1.0,
            filter_selectivity=np.mean([
                o.get('selectivity', 0.5) for o in filter_ops
            ]) if filter_ops else 1.0,
            vectorizable_ops=vectorizable,
            row_level_ops=row_level,
            simd_friendly=simd_friendly,
            estimated_memory_mb=physical_plan.get('estimated_memory_mb', 0),
            spill_risk=physical_plan.get('spill_risk', 0.0)
        )

    def _collect_operators(self, plan: dict) -> List[dict]:
        """递归收集所有算子"""
        ops = [plan]
        for child in plan.get('children', []):
            ops.extend(self._collect_operators(child))
        return ops

3.2 AI代价预测模型

import torch
import torch.nn as nn

class VectorizedCostModel(nn.Module):
    """基于 Transformer 的向量化执行代价预测模型"""

    def __init__(self, feature_dim=16, hidden_dim=128, num_heads=4):
        super().__init__()

        # 算子序列编码器：将计划中的算子序列编码为向量
        self.op_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=hidden_dim,
                nhead=num_heads,
                dim_feedforward=256,
                dropout=0.1,
                batch_first=True
            ),
            num_layers=3
        )

        # 全局特征投影
        self.feature_proj = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim),
            nn.ReLU(),
            nn.LayerNorm(hidden_dim)
        )

        # 多任务输出头
        self.time_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, 64),
            nn.ReLU(),
            nn.Linear(64, 1)  # 预测执行时间（log ms）
        )

        self.cache_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()  # 预测缓存命中率（0-1）
        )

        self.simd_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()  # 预测 SIMD 利用率（0-1）
        )

    def forward(self, op_sequence: torch.Tensor,
                global_features: torch.Tensor):
        """
        op_sequence: [batch, seq_len, hidden_dim] 算子序列
        global_features: [batch, feature_dim] 全局特征
        """
        # 编码算子序列
        op_encoded = self.op_encoder(op_sequence)
        # 取序列均值作为计划级表示
        plan_repr = op_encoded.mean(dim=1)

        # 编码全局特征
        feat_encoded = self.feature_proj(global_features)

        # 拼接两种表示
        combined = torch.cat([plan_repr, feat_encoded], dim=-1)

        # 多任务预测
        pred_time = self.time_head(combined)
        pred_cache = self.cache_head(combined)
        pred_simd = self.simd_head(combined)

        return pred_time, pred_cache, pred_simd


class VectorizedPlanSelector:
    """基于 AI 代价模型的向量化计划选择器"""

    def __init__(self, model: VectorizedCostModel,
                 feature_extractor: PlanFeatureExtractor):
        self.model = model
        self.extractor = feature_extractor
        self.model.eval()

    def select_best_plan(self, candidate_plans: list) -> dict:
        """从候选计划中选择最优向量化执行路径"""
        best_plan = None
        best_score = float('inf')

        for plan in candidate_plans:
            features = self.extractor.extract(plan)
            feature_vec = self._to_tensor(features)

            with torch.no_grad():
                pred_time, pred_cache, pred_simd = self.model(
                    self._plan_to_sequence(plan),
                    feature_vec.unsqueeze(0)
                )

            # 综合评分：执行时间为主，缓存和SIMD为辅
            score = (
                pred_time.item() * 0.6
                - pred_cache.item() * 0.2
                - pred_simd.item() * 0.2
            )

            if score < best_score:
                best_score = score
                best_plan = plan

        return {
            'selected_plan': best_plan,
            'predicted_time_ms': np.exp(pred_time.item()),
            'predicted_cache_hit': pred_cache.item(),
            'predicted_simd_util': pred_simd.item(),
            'score': best_score
        }

四、AI查询优化的局限性与工程权衡

训练数据的冷启动问题：AI 代价模型需要大量"计划→真实性能"的标注数据。新部署的集群缺乏历史数据，模型预测精度低。冷启动阶段只能依赖传统 CBO，待积累足够执行样本后再切换到 AI 模型。通常需要 10000+ 条标注数据才能达到可用精度。

模型推理延迟的叠加：AI 代价模型的推理时间约 5-20ms，对于简单查询（执行时间 < 50ms），优化开销占比过高。生产环境通常设置阈值：仅对预估执行时间 > 100ms 的查询启用 AI 优化，短查询继续使用传统 CBO。

计划空间爆炸：多表 JOIN 的候选计划数量随表数呈阶乘增长（10 表 JOIN 的候选计划超过 10^7 种）。AI 模型无法评估所有候选，需要先通过启发式规则剪枝，保留 Top-K 候选再由 AI 模型排序。剪枝质量直接影响 AI 优化的上限。

分布漂移与在线学习成本：数据分布变化（如大促期间数据量激增）会导致模型预测失准。在线学习可以适应分布变化，但每次更新模型需要重新训练，GPU 资源消耗不可忽视。折中方案是定期批量更新（如每天凌晨），而非实时更新。

五、总结

AI 辅助的向量化查询优化通过学习"计划特征→真实性能"的映射关系，弥补了传统优化器无法感知 CPU 缓存行为和 SIMD 利用率的缺陷。核心架构是"特征提取 + AI 代价预测 + 多任务输出"，预测执行时间、缓存命中率和 SIMD 利用率三个维度。但 AI 优化不是银弹——冷启动依赖大量标注数据、推理延迟对短查询不友好、计划空间需要启发式剪枝、分布漂移需要在线学习。落地建议：冷启动阶段与传统 CBO 并行运行，对比验证 AI 预测精度；仅对预估执行时间 > 100ms 的查询启用 AI 优化；每天凌晨批量更新模型，避免实时训练的资源开销；持续监控 AI 预测误差，误差 > 30% 时回退到传统 CBO。