apache-spark - 按一列中的唯一值随机拆分 DataFrame

Question

我有一个如下所示的 pyspark DataFrame:

+--------+--------+-----------+
| col1   |  col2  |  groupId  |
+--------+--------+-----------+
| val11  | val21  |   0       |
| val12  | val22  |   1       |
| val13  | val23  |   2       |
| val14  | val24  |   0       |
| val15  | val25  |   1       |
| val16  | val26  |   1       |
+--------+--------+-----------+    

每一行都有一个groupId，多行可以有相同的groupId。

我想将此数据随机拆分为两个数据集。但是所有具有特定 groupId 的数据都必须在其中一个拆分中。

这意味着如果 d1.groupId = d2.groupId，则 d1 和 d2 在同一个 split 中。

例如:

# Split 1:

+--------+--------+-----------+
| col1   |  col2  |  groupId  |
+--------+--------+-----------+
| val11  | val21  |   0       |
| val13  | val23  |   2       |
| val14  | val24  |   0       |
+--------+--------+-----------+

# Split 2:
+--------+--------+-----------+
| col1   |  col2  |  groupId  |
+--------+--------+-----------+
| val12  | val22  |   1       |
| val15  | val25  |   1       |
| val16  | val26  |   1       |
+--------+--------+-----------+

在 PySpark 上执行此操作的好方法是什么？我能以某种方式使用 randomSplit 方法吗？

Answer 1

您可以使用 randomSplit 仅拆分不同的 groupId，然后使用结果通过 join 拆分源 DataFrame。

例如:

split1, split2 = df.select("groupId").distinct().randomSplit(weights=[0.5, 0.5], seed=0)
split1.show()
#+-------+
#|groupId|
#+-------+
#|      1|
#+-------+

split2.show()
#+-------+
#|groupId|
#+-------+
#|      0|
#|      2|
#+-------+

现在将这些连接回原始 DataFrame:

df1 = df.join(split1, on="groupId", how="inner")
df2 = df.join(split2, on="groupId", how="inner")

df1.show()
3+-------+-----+-----+
#|groupId| col1| col2|
#+-------+-----+-----+
#|      1|val12|val22|
#|      1|val15|val25|
#|      1|val16|val26|
#+-------+-----+-----+

df2.show()
#+-------+-----+-----+
#|groupId| col1| col2|
#+-------+-----+-----+
#|      0|val11|val21|
#|      0|val14|val24|
#|      2|val13|val23|
#+-------+-----+-----+