python - Spark 中 IF then ELSE 的等价物

Question

我早些时候在这里看到过这个问题，并从中吸取了教训。但是，我不确定为什么当我认为它应该起作用时会出现错误。

我想通过一些规则在现有 Spark DataFrame 中创建一个新列。这是我写的。 iris_spark 是具有分类变量 iris_spark 的数据框，该变量具有三个不同的类别。

from pyspark.sql import functions as F

iris_spark_df = iris_spark.withColumn(
    "Class", 
   F.when(iris_spark.iris_class == 'Iris-setosa', 0, F.when(iris_spark.iris_class == 'Iris-versicolor',1)).otherwise(2))

抛出以下错误。

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-157-21818c7dc060> in <module>()
----> 1 iris_spark_df=iris_spark.withColumn("Class",F.when(iris_spark.iris_class=='Iris-setosa',0,F.when(iris_spark.iris_class=='Iris-versicolor',1)))

TypeError: when() takes exactly 2 arguments (3 given)


---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-157-21818c7dc060> in <module>()
----> 1 iris_spark_df=iris_spark.withColumn("Class",F.when(iris_spark.iris_class=='Iris-setosa',0,F.when(iris_spark.iris_class=='Iris-versicolor',1)))

TypeError: when() takes exactly 2 arguments (3 given)

知道为什么吗？

Answer 1

正确的结构是:

(when(col("iris_class") == 'Iris-setosa', 0)
.when(col("iris_class") == 'Iris-versicolor', 1)
.otherwise(2))

这相当于

CASE 
    WHEN (iris_class = 'Iris-setosa') THEN 0
    WHEN (iris_class = 'Iris-versicolor') THEN 1 
    ELSE 2
END

或者:

(when(col("iris_class") == 'Iris-setosa', 0)
    .otherwise(when(col("iris_class") == 'Iris-versicolor', 1)
        .otherwise(2)))

这相当于:

CASE WHEN (iris_class = 'Iris-setosa') THEN 0 
     ELSE CASE WHEN (iris_class = 'Iris-versicolor') THEN 1 
               ELSE 2 
          END 
END

一般语法:

when(condition, value).when(...)

或

when(condition, value).otherwise(...)

您可能将 Hive IF 条件混淆了:

IF(condition, if-true, if-false)

只能在具有 Hive 支持的原始 SQL 中使用。