Pyspark - データフレームからトレーニングセットとテストセットを作成します

Question

私は以下の写真のようなデータフレームを持っています。私はトレーニングとテストを作成したいと思います。データセットは、CustomerIDとInvoiceNoによって並べ替えられます。顧客ごとに、その顧客の最後の2行を除くすべての行をトレーニングセットとし、各顧客の2番目から最後の行をトレーニングセットにしたいと考えています。

結果は理想的には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")