如何在Polars DataFrame中划分训练集与测试集？现有方法存冗余列问题

这里有两种简洁的方法帮你避免生成多余的split列：

方法1：采样+反连接（推荐） #

直接抽取训练集，再通过反连接得到测试集，全程不需要额外添加列：

# 按8:2比例随机拆分，shuffle=True确保随机性，seed可选用于结果复现
train = df.sample(fraction=0.8, shuffle=True, seed=42)
test = df.join(train, on=df.columns, how="anti")

方法2：优化partition_by写法 #

如果你想保留原来的拆分思路，只需要在partition_by中设置include_key=False，拆分后的数据集就不会包含用于分组的split列：

train, test = (
    df.with_columns(pl.lit(np.random.rand(df.height) > 0.8).alias('split'))
      .partition_by('split', include_key=False)
)

这样处理后，train和test里就不会有多余的split列，省去手动删除的步骤。

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