import pandas as pd
import numpy as np
import random
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
torch.autograd.set_detect_anomaly(True)
# 生成模拟数据库查询日志的函数
def generate_query_log(num_queries=1000):
fields = ["id", "name", "age", "email"]
conditions = ["= ?", "> ?", "< ?", "LIKE ?"]
query_log = pd.DataFrame({
"query_id": range(1, num_queries + 1),
"num_fields": [random.randint(1, len(fields)) for _ in range(num_queries)],
"num_conditions": [random.randint(1, 4) for _ in range(num_queries)],
"use_index": [random.choice([True, False]) for _ in range(num_queries)],
"data_volume": [random.randint(100, 10000) for _ in range(num_queries)] # 模拟的数据量
})
# 计算资源消耗
base_resource_usage_per_field = 0.05 # 字段数对资源消耗的影响较小
base_resource_usage_per_condition = 0.15 # 条件数对资源消耗的影响较大
base_resource_usage_per_volume_unit = 0.001 # 每单位数据量的资源消耗
query_log["resource_usage"] = query_log["num_fields"] * base_resource_usage_per_field + query_log["num_conditions"] * base_resource_usage_per_condition + query_log["data_volume"] * base_resource_usage_per_volume_unit
# 设置执行时间
base_execution_time = 2.0 # 设定一个基础执行时间
execution_time_factor = 0.5 # 用于调整执行时间与资源消耗关系的因子
query_log["execution_time"] = base_execution_time + (query_log["resource_usage"] * execution_time_factor)
query_log["query"] = query_log.apply(lambda row: f"SELECT {', '.join(random.sample(fields, int(row['num_fields'])))} FROM users WHERE {' AND '.join([random.choice(fields) + ' ' + random.choice(conditions) for _ in range(int(row['num_conditions']))])}", axis=1)
return query_log
def extract_features(row):
# 提取查询的长度、涉及的字段数、条件数、是否使用索引、以及数据量
query_length = len(row['query'])
num_fields = row['num_fields']
num_conditions = row['num_conditions']
use_index = 1 if row['use_index'] else 0 # 将布尔值转换为整数
data_volume = row['data_volume'] # 查询涉及的数据量
# 将这些特征组合成一个特征向量
features = [query_length, num_fields, num_conditions, use_index, data_volume]
return features
# 定义策略网络
# 策略网络是一个神经网络，它从状态空间接收输入，并输出每个动作的概率。
# 在我们的例子中，状态空间简化为查询的长度和其中空格的数量（作为查询复杂性的简化代理）
# 而动作空间是一个二元选择：使用索引或不使用索引。
#定义了一个新的类 PolicyNetwork，它继承自 PyTorch 的 nn.Module。
# 在 PyTorch 中，nn.Module 是所有神经网络模型的基类。
# 定义的 PolicyNetwork 是一个特定的神经网络，用于实现策略学习。
class PolicyNetwork(nn.Module):
#初始化函数。它定义了网络的结构。该函数接收两个参数：num_inputs（输入层的维度，即状态空间的大小）
# 和 num_actions（输出层的维度，即动作空间的大小）。
def __init__(self, num_inputs, num_actions):
#调用父类 nn.Module 的初始化函数。在创建新的 PyTorch 模型时，这是标准的做法，它设置了一些背后的基础设施。
# 增加了网络的容量，并添加了一个额外的隐藏层来处理更复杂的输入。这样可以帮助网络学习更复杂的模式，并更好地理解数据库查询的状态。
super(PolicyNetwork, self).__init__()
self.fc1 = nn.Linear(num_inputs, 128) # 定义了网络的第一层，是一个全连接（线性）层。创建了一个从 num_inputs 维到 128 维的线性映射。128 是这一层神经元的数量。
self.fc2 = nn.Linear(128, 64) # 可以添加一个额外的隐藏层
self.fc3 = nn.Linear(64, num_actions) # 第三层，这一层从第二层的 64 个神经元接收输入，并将其映射到 num_actions 个输出，每个输出对应一个动作的概率。
#定义了数据通过网络的前向传播过程，每当你对 PolicyNetwork 的实例进行调用时，这个函数就会被执行。
def forward(self, x):
#在前向传播中，数据首先通过第一层，然后应用ReLU激活函数。ReLU（Rectified Linear Unit）是一个非线性函数，它将所有负值置为0，对正值不做改变。这有助于引入非线性，使网络能够学习更复杂的模式。
x = torch.relu(self.fc1(x)) # 激活函数ReLU
x = torch.relu(self.fc2(x)) # 第二层ReLU激活
#最后，数据通过第三层，并应用softmax函数。softmax函数可以将原始输出转换为概率分布，这些概率总和为1。在这里，它给出了选择每个动作的概率。
return torch.softmax(self.fc3(x), dim=1) # Softmax输出动作概率
# 定义奖励函数
def reward_function(execution_time, resource_usage, max_time, max_resource_usage):
# 标准化执行时间和资源消耗
normalized_time = execution_time / max_time
normalized_resource = resource_usage / max_resource_usage
time_weight = 0.7
resource_weight = 0.3
# 调整奖励计算公式
reward = (1 / normalized_time) * time_weight - normalized_resource * resource_weight
return reward
# 奖励函数是根据执行时间来计算奖励的。在这个简化模型中，奖励是执行时间的倒数，意味着更快的执行时间会获得更高的奖励。
# 在现实情况中，查询的复杂性对于是否应该使用索引有重大影响。一般来说，涉及大量数据的复杂查询可能从索引中受益更多。
def simulate_index_usage(execution_time, resource_usage):
# 假设使用索引会减少执行时间
time_reduction_factor = 0.5 # 可以根据实际情况调整
# 假设使用索引会增加资源消耗
resource_usage_increase_factor = 1.2 # 可以根据实际情况调整
adjusted_execution_time = execution_time * (1 - time_reduction_factor)
adjusted_resource_usage = resource_usage * resource_usage_increase_factor
return adjusted_execution_time, adjusted_resource_usage
def simulate_without_index(execution_time, resource_usage):
# 不使用索引时，执行时间和资源消耗保持不变
return execution_time, resource_usage
# 训练模型
def train(policy_net, optimizer, query_log, num_episodes=1000):
max_execution_time = query_log['execution_time'].max() # 获取最大执行时间用于标准化
max_resource_usage = query_log['resource_usage'].max() # 获取最大资源消耗用于标准化
epsilon = 0.1 # 探索率
for episode in range(num_episodes):
total_reward = 0
for _, row in query_log.iterrows():
# 更新状态表示，包含所有新特征
state = torch.tensor([extract_features(row)], dtype=torch.float32)
if random.random() < epsilon: # 探索
action = torch.tensor([random.choice([0, 1])], dtype=torch.int)
else: # 利用
action_probs = policy_net(state)
m = Categorical(action_probs)
action = m.sample()
# 根据当前的策略网络选择一个动作，然后根据执行时间计算奖励
# execution_time = row['execution_time'] if action.item() == 0 else row['execution_time'] / 2
if action.item() == 0: # 假设动作0代表使用索引
execution_time, resource_usage = simulate_index_usage(row['execution_time'], row['resource_usage'])
else: # 不使用索引
execution_time, resource_usage = simulate_without_index(row['execution_time'], row['resource_usage'])
reward = reward_function(execution_time, resource_usage, max_execution_time, max_resource_usage)
total_reward += reward
# 使用这个奖励更新策略网络，目的是让网络学习如何选择可以获得更高奖励的动作。
optimizer.zero_grad()
# 重新计算log_prob来避免共享计算图
log_prob = Categorical(policy_net(state)).log_prob(action)
loss = -log_prob * reward
loss.backward()
optimizer.step()
print(f'Episode {episode} Total Reward: {total_reward}')
# 生成模拟数据
query_log = generate_query_log()
# 设置DRL环境和模型
num_features = 5 # 这里是状态空间的维度
num_actions = 2 # 二元动作空间：使用索引或不使用索引
policy_net = PolicyNetwork(num_features, num_actions)
optimizer = optim.Adam(policy_net.parameters(), lr=0.01)
# 训练模型
train(policy_net, optimizer, query_log)
