apache-spark - 根据 Spark 中其他列的组值在新列上生成随机值

Question

我们有以下数据框：

------
G1|G2|
1 | 1|
1 | 1|
1 | 2|
2 | 1|
2 | 2|
2 | 3|

所以基于列G1，G2我们有 5 个组：1-1、1-2、2-1、2-2、2-3。

我想创建isSelected具有以下规则的新列：每个组都属于 N 行，我想随机化至少 50% 的行将具有价值1，0否则。每个组必须至少有 1 行，isSelected = 1并且 [1行数] - [0行数] 应小于 1

以下是一个有效的世代：

----------------
G1|G2|isSelected
1 | 1|1
1 | 1|0
1 | 2|1
2 | 1|1
2 | 2|1
2 | 3|1

以下无效：

----------------
G1|G2|isSelected
1 | 1|1
1 | 1|1 --> Not OK, this group has 2 1-row and 0 0-row.
1 | 2|1
2 | 1|1
2 | 2|1
2 | 3|0 --> Not OK, this group has 0 1-row.

如何直接在 Spark 中完成？

标签： apache-sparkrandom

Answer 1

可以使用窗口函数“row_number”为组中的每一行分配序列，并且“sequence/2”余数可用于“isSelected”列：

val allColumnsWindow = Window.partitionBy("G1","G2").orderBy("G2")
df
  .withColumn("rn",row_number().over(allColumnsWindow))
  .withColumn("isSelected", $"rn" % 2)
  .drop("rn")