pandas如何在打乱行、拆分数据集后留存原始索引用于溯源

你当前用的打乱写法存在问题：df[:] = df.sample(frac = 1, random_state = 5).values 只会把采样后的数值按顺序填回原df的固定位置，原索引不会跟随行移动，本质上破坏了原索引和行内容的对应关系，无法满足追溯需求。
正确实现可以按照以下步骤操作，全程保证样本和原始标识绑定，不会错位：

1. 正确打乱数据，绑定原始索引 #

打乱前先把原始索引存为单独的追踪列，后续所有操作都不会丢失对应关系：

import pandas as pd
from sklearn.model_selection import train_test_split

# 将原始索引存为独立列original_id，作为样本追溯的唯一标识
df = df.reset_index(names='original_id')
# 直接对DataFrame采样实现打乱，*不要取values赋值*，此时每行和对应的original_id完全绑定
shuffled_df = df.sample(frac=1, random_state=5).reset_index(drop=True)

2. 按比例拆分数据集 #

所有拆分步骤固定随机种子，保证结果可复现：

# 示例按7:2:1比例拆分训练集、测试集、验证集
# 先拆分出训练集，剩余30%作为临时集待拆分
train_df, temp_df = train_test_split(shuffled_df, test_size=0.3, random_state=5)
# 临时集按2:1拆分出测试集、验证集，对应总占比20%、10%
test_df, dev_df = train_test_split(temp_df, test_size=1/3, random_state=5)

3. 生成索引映射表，清理数据集 #

提取每个样本的所属数据集、在数据集内的位置生成映射表，再删除数据集里的追踪列，得到不保留原始索引的干净数据集：

# 为每个子集生成映射记录：原始id、所属数据集、集内行位置
train_map = train_df[['original_id']].assign(
    dataset='train',
    position=range(len(train_df))
)
test_map = test_df[['original_id']].assign(
    dataset='test',
    position=range(len(test_df))
)
dev_map = dev_df[['original_id']].assign(
    dataset='dev',
    position=range(len(dev_df))
)

# 拼接得到全量索引映射表，可单独存储为文件供后续查询
index_mapping = pd.concat([train_map, test_map, dev_map]).set_index('original_id')

# 删除子集内的追踪列，重置为连续的新索引，得到干净的拆分结果
train_df = train_df.drop(columns=['original_id']).reset_index(drop=True)
test_df = test_df.drop(columns=['original_id']).reset_index(drop=True)
dev_df = dev_df.drop(columns=['original_id']).reset_index(drop=True)

映射表使用方式 #

需要查询某条原始样本的去向时，直接通过原始id查映射表即可，例如查询原始编号100的样本：

# 查询结果会返回样本所属数据集、对应位置
target_info = index_mapping.loc[100]
# 拿到位置后直接从对应数据集取数即可
# 示例：如果样本在测试集，就用test_df.iloc[target_info.position]取数

注意事项 #

• 所有涉及随机采样、拆分的步骤统一固定random_state值，可保证每次运行的打乱、拆分结果完全可复现
• 全程不要对DataFrame取.values或.to_numpy()后再赋值回结构，会丢失列绑定关系，直接操作DataFrame本身即可保证original_id和样本行始终对应
• 映射表可以单独存为csv、pickle格式，后续做误差分析、结果回溯时直接加载使用

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