Как сделать check / try-catch для фрейма данных pyspark?

Question

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

Я попытался обернуть оператор when внутри и после оператора .withColumn

df = df.withColumn('total_new_load',
                     col('existing_load') * (5 - col('tot_reduced_load')))

В основном мне нужно добавить оператор if некоторого вида в синтаксис pyspark, относящийся к моему коду данных, например:

  if tot_reduced_load > 50 
  then 
  tot_reduced_load = 50

Answer 1

Попробуйте это

from pyspark.sql import functions as F
df.withColumn("tot_reduced_load ", F.when(F.col("tot_reduced_load")>50,50)).otherwise(F.col("tot_reduced_load"))

Answer 2

Попробуйте это -

Пример данных:

df = spark.createDataFrame([(1,30),(2,40),(3,60)],['row_id','tot_reduced_load'])

df.show()

#+------+----------------+
#|row_id|tot_reduced_load|
#+------+----------------+
#|     1|              30|
#|     2|              40|
#|     3|              60|
#+------+----------------+

Опция 1: withColumn

from pyspark.sql import functions as psf

tot_reduced_load_new  = psf.when(psf.col("tot_reduced_load") > 50 , 50).otherwise(psf.col("tot_reduced_load"))

df.withColumn("tot_reduced_load_new",tot_reduced_load_new ).show()

#+------+----------------+--------------------+
#|row_id|tot_reduced_load|tot_reduced_load_new|
#+------+----------------+--------------------+
#|     1|              30|                  30|
#|     2|              40|                  40|
#|     3|              60|                  50|
#+------+----------------+--------------------+

Опция 2: selectExpr

df.selectExpr("*","CASE WHEN tot_reduced_load > 50 THEN 50 ELSE tot_reduced_load END AS tot_reduced_load_new").show()

#+------+----------------+--------------------+
#|row_id|tot_reduced_load|tot_reduced_load_new|
#+------+----------------+--------------------+
#|     1|              30|                  30|
#|     2|              40|                  40|
#|     3|              60|                  50|
#+------+----------------+--------------------+