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