Заменить определенную группу в столбце данных pyspark случайным элементом из списка

Question

Допустим, датафрейм выглядит следующим образом:

ls = [
    ['1', -9.78],
    ['2', 5.38],
    ['1', 8.86],
    ['2', -0.47],
    ['1', -0.19],
    ['1', 4.78],
    ['1', -9.23],
    ['2', -89.32]
]
test = spark.createDataFrame(pd.DataFrame(ls, columns=['col1', 'col2']))
test.show()

выход:

+----+------+
|col1|  col2|
+----+------+
|   1| -9.78|
|   2|  5.38|
|   1|  8.86|
|   2| -0.47|
|   1| -0.19|
|   1|  4.78|
|   1| -9.23|
|   2|-89.32|
+----+------+

Я хочу заменить всю строку, в которой значение в столбце col1 == 1, случайным выбором из списка элементов: ['a', 'b', 'c'] (с заменой).

Например, результат будет выглядеть так:

+----+------+
|col1|  col2|
+----+------+
|   a| -9.78|
|   2|  5.38|
|   a|  8.86|
|   2| -0.47|
|   c| -0.19|
|   b|  4.78|
|   a| -9.23|
|   2|-89.32|
+----+------+

Я новичок в Pyspark, и я совершенно не понимаю, когда использовать такие операции, как когда (), где (), withColumn (), select () и т. Д. Я искал на этом сайте, но не могу ' Я не могу найти ответ на мой вопрос, поэтому я очень надеюсь, что кто-нибудь сможет помочь!

Answer 1

Вы можете определить функцию для случайного выбора значения из вашего списка.Используя pyspark.sql.functions.rand(), вы можете сгенерировать равномерное случайное число от 0 до 1. На основе значения выбранного случайного значения вы можете выбрать индекс из списка.

Например, в вашем случае, когда в вашем списке 3 элемента:

• Выберите 'a', если случайное число меньше 1/3
• Выберите 'b', если случайное число меньше 2/3
• В противном случае выберите 'c'

Вот функция, которая может быть использована для обобщения этого для любой длинысписок выбора:

from pyspark.sql.functions import rand, when

def pickRandomValue():
    picklist = ['a', 'b', 'c']
    N = len(picklist)
    r = rand()
    c = when(r <= 1.0/N, picklist[0])
    for i in range(1, N-1):
        c = c.when(r <= (i+1.0)/N, picklist[i])
    c = c.otherwise(picklist[-1])
    return c

Теперь используйте when для изменения col1:

from pyspark.sql.functions import col

test = test.withColumn(
    "col1", 
    when(
        col("col1") == 1,
        pickRandomValue()
    ).otherwise(col("col1").cast("string"))
)
test.show()
#+----+------+
#|col1|  col2|
#+----+------+
#|   c| -9.78|
#|   2|  5.38|
#|   c|  8.86|
#|   2| -0.47|
#|   b| -0.19|
#|   a|  4.78|
#|   a| -9.23|
#|   2|-89.32|
#+----+------+

Однако я должен повторить использование "спредостережение "предупреждение о случайных числах в искре.Поскольку искра ленива, вызов rand() будет повторяться при каждом действии.Чтобы продемонстрировать это, вот результат второго вызова show():

test.show()
#+----+------+
#|col1|  col2|
#+----+------+
#|   c| -9.78|
#|   2|  5.38|
#|   b|  8.86|
#|   2| -0.47|
#|   b| -0.19|
#|   b|  4.78|
#|   b| -9.23|
#|   2|-89.32|
#+----+------+

Как видите, результаты разные.