使用PySpark将2038行DataFrame拆分为2008行训练集与30行测试集

如果需要将2038行的PySpark DataFrame精确拆分为2008行训练集和30行测试集，有两种可行方案：

方案一：精确按行数拆分 #

按原有顺序拆分 #

直接提取前2008行作为训练集，剩余行自动成为测试集：

# 假设原DataFrame变量名为df
train_df = df.limit(2008)
test_df = df.exceptAll(train_df)

limit(2008)会获取DataFrame的前2008行，exceptAll会返回原DataFrame中不在训练集内的所有行，刚好得到30行测试集。

随机打乱后拆分 #

如果需要训练集和测试集是随机分布的，先打乱原DataFrame的顺序再拆分：

from pyspark.sql.functions import rand

# 随机打乱DataFrame行顺序
shuffled_df = df.orderBy(rand())
# 提取训练集
train_df = shuffled_df.limit(2008)
# 得到测试集
test_df = shuffled_df.exceptAll(train_df)

方案二：按比例拆分（无法保证精确行数） #

如果对行数精度要求不高，可以用randomSplit按比例拆分，比例对应2008/2038≈0.985和30/2038≈0.015：

train_df, test_df = df.randomSplit([0.985, 0.015], seed=42)

注意：这种方法是按随机比例分配，实际行数可能和目标有微小偏差，适合不需要绝对精确行数的场景。

内容的提问来源于stack exchange，提问作者Lokeesh Bharathi S