Создать случайное значение для нового столбца, основываясь на значении группы других столбцов в Spark

Question

У нас есть следующий Dataframe:

------
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?

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")