Pyspark - 从数据框创建训练集和测试集

2024-05-31 23:59:48 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个像下面这张照片的数据帧。我想创建一个培训和测试设置。数据集按CustomerID和InvoiceNo排序。对于每个客户,除了最后2行之外,我希望将该客户的每一行作为培训集,而每个客户的第二到最后一行将成为一个培训集。在

理想的结果是1个大型训练集和1个测试集。有没有一种有效的方法来使用PySpark?提前谢谢你的帮助

enter image description here


Tags: 数据方法客户排序照片pyspark理想customerid
1条回答
网友
1楼 · 发布于 2024-05-31 23:59:48

您可以随时添加一个索引,并根据该索引进行筛选,不确定是否有比这更有效的方法。在

from pyspark.sql.window import Window
from pyspark.sql import functions as func

window = Window.partitionBy(func.col("CustomerID"))\
        .orderBy(func.col("InvoiceNo").desc())
df = df.select('*', func.rank().over(window).alias('rank'))

train = df.filter("rank > 2")
test = df.filter("rank <= 2")

相关问题 更多 >