通过保留顺序，根据 id 列将 Spark DataFrame 拆分为两个 DataFrame（70% 和 30%）

Posted 2023-04-18

【中文标题】通过保留顺序，根据 id 列将 Spark DataFrame 拆分为两个 DataFrame（70% 和 30%）

【英文标题】：Split Spark DataFrame into two DataFrames (70% and 30% ) based on id column by preserving order

【发布时间】：2018-10-23 21:41:43

【问题描述】：

我有一个类似的 spark 数据框

id  start_time   feature
1   01-01-2018   3.567
1   01-02-2018   4.454
1   01-03-2018   6.455
2   01-02-2018   343.4
2   01-08-2018   45.4
3   02-04-2018   43.56
3   02-07-2018   34.56
3   03-07-2018   23.6

我希望能够根据 id 列将其拆分为两个数据帧。所以我应该按 id 列分组，按 start_time 排序并将 70% 的行放入一个数据帧和 30通过保留顺序将 % 的行放入另一个数据框中。结果应如下所示：

Dataframe1:
id  start_time   feature
1   01-01-2018   3.567
1   01-02-2018   4.454
2   01-02-2018   343.4
3   02-04-2018   43.56
3   02-07-2018   34.56

Dataframe2:
1   01-03-2018   6.455
2   01-08-2018   45.4
3   03-07-2018   23.6

我正在使用带有 python 的 Spark 2.0。实现这一点的最佳方法是什么？

【问题讨论】：

Stratified sampling in Spark 和 Stratified sampling with pyspark 的可能重复

我已经检查了该答案，但它没有回答如何在拆分时保留订单

从您发布的示例中，您似乎将每个组中的最新日期作为拆分 dfs 的一部分，这是一个要求吗？如果不是，我会说只是排序，然后 start_time 然后随机拆分

【参考方案1】：

我必须这样做是创建两个窗口：

w1 =  Window.partitionBy(df.id).orderBy(df.start_time)
w2 =  Window.partitionBy(df.id)

df = df.withColumn("row_number",F.row_number().over(w1))\
                     .withColumn("count",F.count("id").over(w2))\
                     .withColumn("percent",(F.col("row_number")/F.col("count")))
train = df.filter(df.percent<=0.70)
test = df.filter(df.percent>0.70)