python - PySpark - 条件语句

Question

我是 PySpark 的新手，想知道您是否可以指导我如何将以下 SAS 代码转换为 PySpark。

SAS 代码：

If ColA > Then Do;
    If ColB Not In ('B') and ColC <= 0 Then Do;
         New_Col = Sum(ColA, ColR, ColP);
    End;
    Else Do;
         New_Col = Sum(ColA + ColR)
    End;
End;
Else Do;
   If ColB Not in ('B') and ColC <= 0 then do;
     New_Col = Sum(ColR, ColP);
   end;
   Else Do;
     New_Col = ColR;
   End;
End;

目前，以下是我正在使用的 PySpark 逻辑：

df.withColumn('New_Col', when(ColA > 0 & ColB.isin(['B']) == False & ColC <= 0, col('ColA') + Col('ColR') + Col('ColP'))
...
...

这是最优化的方法还是有更好的代码方法？

感谢您的指导！

标签： pythonapache-sparkpyspark

Answer 1

您的代码符合需要，但条件应包含在括号内

from pyspark.sql import functions as F

(df
    .withColumn('New_Col', F
        .when((F.col('ColA') > 0) & (F.col('ColB').isin(['B']) == False) & (F.col('ColC') <= 0), F.col('ColA') + F.Col('ColR') + F.Col('ColP'))
    )
)