Примените формулу в pyspark

#python #pyspark

Вопрос:

Я новичок в spark, и у меня есть некоторые сомнения по поводу работы с df.
Моя проблема в том, что мне нужно применить формулу к столбцу df pyspark, используя значения из других столбцов.
У меня есть следующие df

  ------- ------- ------- ------- ------- -------  | count1| count2|val__00|val__01|val__02|val__03|  ------- ------- ------- ------- ------- -------  | 3| 1| 1.0| 0.0| 8.0| 0.0| | 4| 2| 0.0| 1.379| 1.49| 1.373| | 2| 5| 0.7| 0.0| 70.4| 0.0| | 1| 8| 0.0| 4.0| 0.0| 0.0|  ------- ------- ------- ------- ------- -------    

Мне нужно применить следующую формулу к столбцам val__xx для каждой строки:

 val__xx = val__xx   (count1*count2)   

таким образом, окончательный df будет

  ------- ------- ------- ------- ------- -------  | count1| count2|val__00|val__01|val__02|val__03|  ------- ------- ------- ------- ------- -------  | 3| 1| 4.0| 3.0| 11.0| 3.0| | 4| 2| 8.0| 9.379| 9.49| 9.373| | 2| 5| 10.7| 10.0| 80.4| 10.0| | 1| 8| 8.0| 12.0| 8.0| 8.0|  ------- ------- ------- ------- ------- -------   

Я думаю применить функцию udf, но я не знаю, как передать более одного столбца. Можно ли создать функцию для передачи более одного столбца?
Я реализовал код ниже, но я не знаю, как передать cols val__xx

 def calculate(c, count1, count2):  return c   (count1*count2)   calculateUDF = udf(lambda x: calculate(x, count1, count2))  df_final = df.apply(calculateUDF(col(val__xx????), col(count1), col(count2))  

Ответ №1:

Вы можете использовать withColumn в цикле for, нет необходимости в udf.

 from pyspark.sql import functions as f  for i in range(4):  df = df.withColumn(f'val__0{i}', f.col('count1') * f.col('count2')   f.col(f'val__0{i}'))  df.show()   ------ ------ ------- ------- ------- -------   |count1|count2|val__00|val__01|val__02|val__03|  ------ ------ ------- ------- ------- -------  | 3| 1| 4.0| 3.0| 11.0| 3.0| | 4| 2| 8.0| 9.379| 9.49| 9.373| | 2| 5| 10.7| 10.0| 80.4| 10.0| | 1| 8| 8.0| 12.0| 8.0| 8.0|  ------ ------ ------- ------- ------- -------   

Если ваши столбцы «значение» достигают двузначных цифр, вам нужно i будет заполнить их нулями слева.

Комментарии:

1. большое спасибо за ваш ответ! Это работает!