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

Question

我有一个如下图所示的数据框。我想用它创建一个训练和测试集。数据集按 CustomerID 和 InvoiceNo 排序。对于每个客户，我想将该客户除了最后两行之外的每一行作为训练集，而每个客户的倒数第二行将成为训练集。

理想情况下，结果是 1 个巨大的训练集和 1 个测试集。有没有一种有效的方法可以使用 PySpark 来做到这一点？非常感谢您提前提供的帮助

Answer 1

您始终可以添加索引和基于该索引的过滤器 - 不确定是否有比这更有效的方法。

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")