Случайно разделить DataFrame по уникальным значениям в одном столбце

Question

У меня есть фрейм данных pyspark, подобный следующему:

+--------+--------+-----------+
| 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 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 с, а затем использовать результаты, чтобы разделить исходный DataFrame, используя join.

Например:

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

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

Теперь присоедините их обратно к исходному фрейму данных:

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|
#+-------+-----+-----+