PySpark中np.where的等效实现方法

要实现你这段Pandas代码的功能，PySpark里有两种常用的方式，逻辑和np.where完全对应，我给你详细拆解一下：

方法1：使用when/otherwise函数（最贴近Pandas的写法） #

这是PySpark中最常用的条件列生成方式，语法逻辑和np.where几乎一致：

from pyspark.sql import SparkSession
from pyspark.sql.functions import when

# 初始化SparkSession（PySpark必须的启动步骤）
spark = SparkSession.builder.appName("pandas_to_pyspark").getOrCreate()

# 创建和Pandas结构一致的DataFrame
data = [("A", "Z"), ("B", "Z"), ("B", "X"), ("C", "Y")]
df = spark.createDataFrame(data, schema=["Type", "Set"])

# 添加color列，完全等效于Pandas里的np.where逻辑
df = df.withColumn("color", when(df["Set"] == "Z", "green").otherwise("red"))

# 查看结果
df.show()

方法2：使用SQL表达式（适合熟悉SQL的开发者） #

如果你更习惯SQL的CASE WHEN语法，也可以用expr函数直接写SQL风格的逻辑：

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

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

data = [("A", "Z"), ("B", "Z"), ("B", "X"), ("C", "Y")]
df = spark.createDataFrame(data, schema=["Type", "Set"])

# 用SQL的CASE WHEN实现同样的条件判断
df = df.withColumn("color", expr("CASE WHEN Set = 'Z' THEN 'green' ELSE 'red' END"))

df.show()

输出结果 #

两种方法都会得到和Pandas完全一致的结果：

+----+---+-----+
|Type|Set|color|
+----+---+-----+
|   A|  Z|green|
|   B|  Z|green|
|   B|  X|  red|
|   C|  Y|  red|
+----+---+-----+

小提醒 #

PySpark的DataFrame是不可变的，所以每次调用withColumn后都要重新赋值给变量哦～如果有更复杂的多分支条件，还可以链式调用多个when来扩展，比如when(条件1, 值1).when(条件2, 值2).otherwise(默认值)

内容的提问来源于stack exchange，提问作者adil blanco