我有一个如图所示的数据框。我想从中创建一个训练集和一个测试集。数据集是按客户ID和发票号排序的。对于每个客户，我希望将除了最后两行之外的所有行作为训练集，而每个客户的倒数第二行将成为测试集的一部分。

理想的结果将是一个巨大的训练集和一个测试集。使用PySpark有没有高效的方法来做到这一点？非常感谢您的帮助。

回答：

你总是可以添加一个索引并根据该索引进行过滤——我不确定是否有比这更有效的方法。

from pyspark.sql.window import Windowfrom pyspark.sql import functions as funcwindow = 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")