在 PySpark Dataframe 中按条件计算值

在 PySpark 中，可以使用 Dataframe 对象处理结构化数据。有时候，需要在 Dataframe 中按条件计算某些值，本文将介绍如何使用 PySpark 在 Dataframe 中按条件计算值。

创建示例 Dataframe

首先，我们需要创建一个示例 Dataframe：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

spark = SparkSession.builder.appName("example").getOrCreate()

# 创建 schema
schema = StructType([
    StructField("name", StringType(), True),
    StructField("age", IntegerType(), True),
    StructField("gender", StringType(), True)
])

# 创建 Dataframe
data = [("Alice", 25, "female"),
        ("Bob", 30, "male"),
        ("Charlie", 35, "male"),
        ("David", 40, "male"),
        ("Ella", 45, "female")]

df = spark.createDataFrame(data, schema=schema)

df.show()

输出：

+-------+---+------+
|   name|age|gender|
+-------+---+------+
|  Alice| 25|female|
|    Bob| 30|  male|
|Charlie| 35|  male|
|  David| 40|  male|
|   Ella| 45|female|
+-------+---+------+

按条件计算值

使用 when() 和 otherwise()

有时候需要根据某个条件计算新的列，可以使用 PySpark 的 when() 和 otherwise() 函数。

例如，我们可以根据 age 列的值计算一个新列 group，表示根据年龄分组。年龄小于 30 岁的为 young，大于等于 30 岁小于 40 岁的为 middle，大于等于 40 岁的为 old。

from pyspark.sql.functions import when

df_with_group = df.withColumn("group", 
                              when(df["age"] < 30, "young")
                              .when((df["age"] >= 30) & (df["age"] < 40), "middle")
                              .otherwise("old"))

df_with_group.show()

输出：

+-------+---+------+------+
|   name|age|gender| group|
+-------+---+------+------+
|  Alice| 25|female| young|
|    Bob| 30|  male|middle|
|Charlie| 35|  male|middle|
|  David| 40|  male|   old|
|   Ella| 45|female|   old|
+-------+---+------+------+

使用 expr()

除了 when() 和 otherwise() 函数，还可以使用 PySpark 的 expr() 函数来计算新列。

例如，我们可以根据 gender 列的值计算一个新列 label，表示根据性别分组。如果 gender 为 female 则 label 为 0，否则为 1。

from pyspark.sql.functions import expr

df_with_label = df.withColumn("label", expr("CASE WHEN gender = 'female' THEN 0 ELSE 1 END"))

df_with_label.show()

输出：

+-------+---+------+-----+
|   name|age|gender|label|
+-------+---+------+-----+
|  Alice| 25|female|    0|
|    Bob| 30|  male|    1|
|Charlie| 35|  male|    1|
|  David| 40|  male|    1|
|   Ella| 45|female|    0|
+-------+---+------+-----+

总结

本文介绍了如何使用 PySpark 在 Dataframe 中按条件计算值。当需要根据某个条件计算新的列时，可以使用 when() 和 otherwise() 或 expr() 函数。