使用scikit-learn train_test_split拆分后y_test出现NaN的问题排查与解决

样本数据 #

customer_id   revenue_m10   revenue_m9   revenue_m8  target
   1             1234         1231        1256         1239
   2             5678         3425        3255         2345

所用代码 #

X_train,X_test,y_train, y_test  = train_test_split(
    sample_set_df[all_features], 
    sample_set_df[target_var], 
    test_size=0.3
)

问题排查与解决 #

先排查最常见原因：原始数据集存在缺失值 #

y_test出现NaN的首要原因大概率是原始数据集的target列本身就存在缺失值，拆分后这些缺失值被分到了测试集里。可以先执行以下代码验证：

# 检查target列的缺失值数量
print(sample_set_df[target_var].isna().sum())

如果确实存在缺失值，可按场景选择处理方式：

• 缺失值占比极低：直接删除对应行

  sample_set_df = sample_set_df.dropna(subset=[target_var])
  

• 缺失值占比高：根据业务场景选择填充方式（如均值、中位数）

  # 用中位数填充示例
  sample_set_df[target_var] = sample_set_df[target_var].fillna(sample_set_df[target_var].median())
  

若为索引问题的解决方法 #

如果原始数据target列无缺失，再考虑索引不匹配的情况：

1. 重置数据集索引，确保X和y的索引完全对齐

   sample_set_df = sample_set_df.reset_index(drop=True)
   

2. 确认all_features和target_var对应的列都存在于sample_set_df中，且没有因之前的数据操作（如筛选、切片）导致的索引断层。

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