基于图神经网络的数据库索引推荐：工作负载感知的自动索引

国医中兴

160人浏览 · 2026-06-10 23:51:13

国医中兴 · 2026-06-10 23:51:13 发布

基于图神经网络的数据库索引推荐：工作负载感知的自动索引

cover

一、索引推荐的搜索空间爆炸：从人工经验到智能搜索

数据库索引是查询性能的关键因素，但选择哪些列建索引是一个组合优化问题。一张 20 列的表，最多 3 列联合索引的搜索空间超过 7000 种。全库 100 张表的搜索空间达到 7 亿种，人工经验根本无法穷举。

生产环境中，索引推荐面临三个核心痛点：第一，搜索空间巨大——列组合的排列组合导致候选索引数量指数级增长；第二，索引间的相互影响——新建索引可能使其他索引失效，需要全局优化而非局部贪心；第三，写入代价的权衡——每个索引都会增加写入开销，索引数量需要在查询加速和写入代价之间平衡。

这个问题的本质是：索引推荐是一个组合优化问题——在候选索引空间中找到使查询总代价最小、且写入代价可接受的索引子集。图神经网络（GNN）可以将查询、表和列建模为图结构，学习它们之间的依赖关系，高效搜索最优索引组合。

二、GNN 索引推荐的底层机制

flowchart TB
    subgraph 图构建["工作负载图构建"]
        Q1[查询1] --> C1[列A]
        Q1 --> C2[列B]
        Q2[查询2] --> C2
        Q2 --> C3[列C]
        C1 --> T1[表1]
        C2 --> T1
        C3 --> T1
    end

    subgraph GNN推理["GNN 推理"]
        GRAPH[工作负载图] --> EMBED[节点嵌入]
        EMBED --> SCORE[索引评分<br/>每个候选索引的收益分数]
    end

    SCORE --> SELECT[索引选择<br/>约束优化]
    SELECT --> |写入代价约束| RESULT[推荐索引集]

关键机制解析：

工作负载图：将查询工作负载建模为异构图——查询节点、列节点和表节点，边表示查询访问列、列属于表的关系。GNN 通过消息传递学习节点的嵌入表示。
索引评分：GNN 输出每个候选索引的收益分数——考虑查询加速效果和写入代价。分数高的索引优先推荐。
约束优化：在写入代价预算约束下，选择总收益最大的索引子集。这是一个 0-1 背包问题，可以用贪心或动态规划求解。

三、GNN 索引推荐的实现

3.1 工作负载图构建

import torch
from torch_geometric.data import HeteroData

class WorkloadGraphBuilder:
    """工作负载图构建器"""

    def build(self, workload: list, schema: dict) -> HeteroData:
        data = HeteroData()

        # 查询节点特征: [执行频率, 平均耗时, 返回行数]
        query_features = []
        query_id_map = {}

        for i, query in enumerate(workload):
            query_id_map[query.id] = i
            query_features.append([
                query.frequency,
                query.avg_latency_ms,
                query.avg_rows,
            ])

        data["query"].x = torch.tensor(query_features, dtype=torch.float)

        # 列节点特征: [基数, 数据类型编码, 是否已有索引]
        column_features = []
        column_id_map = {}

        for i, col in enumerate(schema.columns):
            column_id_map[col.full_name] = i
            column_features.append([
                col.cardinality,
                col.dtype_encoded,
                1.0 if col.has_index else 0.0,
            ])

        data["column"].x = torch.tensor(column_features, dtype=torch.float)

        # 查询→列边: 查询访问了哪些列
        edge_src, edge_dst = [], []
        for query in workload:
            for col_name in query.accessed_columns:
                edge_src.append(query_id_map[query.id])
                edge_dst.append(column_id_map[col_name])

        data["query", "accesses", "column"].edge_index = torch.tensor(
            [edge_src, edge_dst], dtype=torch.long
        )

        # 列→表边
        edge_src, edge_dst = [], []
        for i, col in enumerate(schema.columns):
            edge_src.append(i)
            edge_dst.append(col.table_id)

        data["column", "belongs_to", "table"].edge_index = torch.tensor(
            [edge_src, edge_dst], dtype=torch.long
        )

        return data

3.2 GNN 索引评分模型

import torch.nn as nn
from torch_geometric.nn import HeteroConv, SAGEConv

class IndexRecommendationGNN(nn.Module):
    """索引推荐GNN模型"""

    def __init__(self, hidden_dim=128):
        super().__init__()

        self.conv1 = HeteroConv({
            ("query", "accesses", "column"): SAGEConv((-1, -1), hidden_dim),
            ("column", "belongs_to", "table"): SAGEConv((-1, -1), hidden_dim),
        })

        self.conv2 = HeteroConv({
            ("query", "accesses", "column"): SAGEConv((-1, -1), hidden_dim),
            ("column", "belongs_to", "table"): SAGEConv((-1, -1), hidden_dim),
        })

        # 索引评分头
        self.index_scorer = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, data):
        # 两层消息传递
        x_dict = self.conv1(data.x_dict, data.edge_index_dict)
        x_dict = {k: v.relu() for k, v in x_dict.items()}
        x_dict = self.conv2(x_dict, data.edge_index_dict)

        return x_dict

    def score_index(
        self, column_embeddings: torch.Tensor, query_embeddings: torch.Tensor
    ) -> torch.Tensor:
        """为候选索引评分"""
        # 拼接列嵌入和查询嵌入
        combined = torch.cat([
            column_embeddings.unsqueeze(1).expand(-1, query_embeddings.size(0), -1),
            query_embeddings.unsqueeze(0).expand(column_embeddings.size(0), -1, -1),
        ], dim=-1)

        return self.index_scorer(combined).squeeze(-1)