Vectorization.py

import torch

# 递归构建查询计划树的特征向量
def construct_feature_vector(query_plan_tree):
    def encode_sub_plan(sub_plan):
        encoded = []
        node_type_id = int(sub_plan.get("Node Type ID", 0))
        encoded.append(node_type_id)
        encoded.append((sub_plan.get("Total Cost", 0) / 1000000))
        encoded.append((sub_plan.get("Plan Rows", 0)))

        for sub_sub_plan in sub_plan.get("Plans", []):
            encoded.extend(encode_sub_plan(sub_sub_plan))
        return encoded

    encoded_sequence = []
    encoded_sequence.extend(encode_sub_plan(query_plan_tree))

    MAX_LENGTH = 30  # 例如
    if len(encoded_sequence) < MAX_LENGTH:
        encoded_sequence += [0] * (MAX_LENGTH - len(encoded_sequence))
    elif len(encoded_sequence) > MAX_LENGTH:
        encoded_sequence = encoded_sequence[:MAX_LENGTH]
    return torch.tensor(encoded_sequence, dtype=torch.float32)