将列中的不同值用于并行Pyspark数据帧以随机选择值

Question

我是Pyspark的新手，我遇到了一些小问题。假设我有两个字段：

| firstName | lastName |
|-----------|----------|
| Bill      | Apple    |
| Mike      | Apple    |
| Jeff      | Apple    |
| Paul      | Apple    |
| George    | Bowers   |
| Kevin     | Bowers   |
| Leon      | Bowers   |
| Fred      | Bowers   |

我的问题是如何为姓氏的每个不同值随机选择2行？像这样：

| firstName | lastName |
|-----------|----------|
| Jeff      | Apple    |
| Bill      | Apple    |
| Fred      | Bowers   |
| Kevin     | Bowers   |

我当时的想法是生成一个不同姓氏的列表并运行for循环，但显然不推荐使用Pyspark的框架。我认为在这种情况下，使用并行计算是推荐的方法

Answer 1

你可以用一些analytic functions魔法来做到这一点✨

from pyspark.sql import functions as F, Window

df.withColumn(
    "r", F.row_number().over(Window.partitionBy("lastName").orderBy(F.rand()))
).where(F.col("r") <= 2).drop("r").show()
+    -+    +
|firstName|lastName|
+    -+    +
|   Paul  |  Apple |
|   Bill  |  Apple |
|   Kevin |  Bowers|
|   Leon  |  Bowers|
+    -+    +

例如，如果我重新运行它：

+    -+    +
|firstName|lastName|
+    -+    +
|   Paul  |  Apple |
|   Mike  |  Apple |
|   Fred  |  Bowers|
|   Kevin |  Bowers|
+    -+    +