pandas拆分DataFrame为训练/测试集未返回正确DataFrame问题求助

问题复现 #

你当前使用的拆分代码如下：

train_prices=int(len(prices)*0.65)
test_prices=len(prices)-train_prices
train_data,test_data=prices.iloc[0:train_prices,:].values,prices.iloc[train_prices:len(prices),:1].values

实际运行后未得到预期的可用于ML建模的DataFrame格式数据集，仅输出两个数值：

11156746, 813724

对应标识为training_size,test_size，无法支撑后续建模工作。

错误原因 #

• 输出对象错位：你实际打印输出的是train_prices、test_prices两个存储数据集长度的整数变量，不是拆分得到的train_data、test_data数据集对象
• 多余的类型转换：代码末尾调用.values会把pandas DataFrame转为numpy数组，丢失DataFrame的列名、索引、数据类型等元信息，不符合你要获取DataFrame的需求
• 测试集切片逻辑错误：测试集切片参数写为:1，代表仅截取原DataFrame的第1列，会导致训练集和测试集列数不一致，后续建模直接报维度匹配错误

修复代码 #

直接去掉.values调用，修正测试集切片范围即可得到正确的DataFrame格式训练、测试集：

# 计算拆分切分点
train_size = int(len(prices) * 0.65)
# 按顺序拆分，保留完整DataFrame结构
train_data = prices.iloc[:train_size, :]
test_data = prices.iloc[train_size:, :]

校验方法 #

拆分完成后执行以下代码验证结果正确性：

# 输出两个数据集的维度
print(f"训练集维度：{train_data.shape}")
print(f"测试集维度：{test_data.shape}")
# 校验训练集、测试集列数一致
assert train_data.shape[1] == test_data.shape[1], "训练集与测试集列数不匹配"

补充说明 #

• 上述按顺序拆分的逻辑适用于时间序列类预测任务；如果是普通监督学习任务，建议使用随机分层拆分避免数据分布偏移
• 若后续模型训练需要numpy数组格式输入，可在拆分完成后对对应DataFrame调用.to_numpy()方法单独转换，不要在拆分环节直接转换，避免丢失元信息增加调试成本

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