如何正确预处理Dask DataFrame数据并解决ML训练拆分时y维度异常

问题根源 #

你代码里的核心错误是覆盖了正确的y变量：

• 一开始你已经通过dd.get_dummies(df['attack_map'].to_frame().categorize()).values生成了one-hot编码后的7维y数组，但后面又执行了y = df.to_dask_array(lengths=True)，这会把整个数据集（包含42列特征+标签）转成数组，导致y和X形状完全一致，丢失了one-hot编码的结果。
• 另外，直接用sklearn的train_test_split处理Dask数组可能存在兼容性问题，建议使用dask_ml.model_selection.train_test_split，它专门适配Dask的分布式数据结构。

修正后的可运行代码 #

import dask.dataframe as dd
from dask_ml.model_selection import train_test_split

# 读取数据集
df = dd.read_hdf(root_folder+"hdf/"+hdf_name, hdf_name.split(".")[0])

def encode_numeric_zscore(df, name, mean=None, standard_deviation=None):
    if mean is None:
        mean = df[name].mean()
    if standard_deviation is None:
        standard_deviation = df[name].std()
    df[name] = (df[name] - mean) / standard_deviation

# 对特征列做标准化
for column in df.columns:
    if column != 'attack_map':
        encode_numeric_zscore(df, column) 

# 分离特征和标签
X_columns = df.columns.drop('attack_map')
X = df[X_columns].to_dask_array(lengths=True)  # 直接转特征列为dask array
# 对标签做one-hot编码并转为dask array
y = dd.get_dummies(df['attack_map'].categorize()).to_dask_array(lengths=True)

# 确认类型和形状（可选）
print("X类型:", type(X))
print("y类型:", type(y))
print("X形状:", X.shape)
print("y形状:", y.shape)

# 使用dask_ml的train_test_split拆分数据
X_train, X_val, y_train, y_val = train_test_split(
    X, y, test_size=0.2, shuffle=True, random_state=2
) 

# 打印拆分后的形状（需compute触发计算）
print("X_train形状:", X_train.shape.compute())
print("y_train形状:", y_train.shape.compute())
print("X_val形状:", X_val.shape.compute())
print("y_val形状:", y_val.shape.compute())

关键修正点说明 #

• 保留正确的y变量：删除了y = df.to_dask_array(lengths=True)这一行，改用one-hot编码结果生成dask array，确保y是7维的分类标签数组。
• 使用dask_ml拆分工具：dask_ml.model_selection.train_test_split支持分布式数据的高效拆分，避免sklearn原生方法的兼容性问题。
• 统一数组转换方式：用to_dask_array(lengths=True)替代.values，明确生成带长度信息的Dask数组，保障后续操作稳定性。
• 形状打印需compute：Dask数组的形状是延迟计算的，必须调用.compute()才能获取实际数值。

内容的提问来源于stack exchange，提问作者Nootey