如何对pyspark数据框进行检查/尝试捕获？

Question

我有一个数据框，该数据框基于对现有列的缩减计算来创建新列。 我需要检查一下，如果所使用的减少值高于特定阈值，则应使其等于阈值/应该不超过阈值。

我尝试在.withColumn语句之内和之后包装一个when语句

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

基本上，我需要在与我的数据帧代码相关的pyspark语法中添加某种if语句，例如：

  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|
#+------+----------------+--------------------+