PySpark如何根据列表列是否含指定值生成二进制列

实现步骤 #

你可以用PySpark内置的array_contains函数结合类型转换，快速生成每个颜色的二进制标识列：

1. 初始化环境并创建示例DataFrame

from pyspark.sql import SparkSession
from pyspark.sql.functions import array_contains, col

# 启动Spark会话
spark = SparkSession.builder.appName("ColorBinaryFlags").getOrCreate()

# 构建示例数据
sample_data = [
    (["Red","Blue","Green","Black"],),
    (["Blue","Green"],)
]
df = spark.createDataFrame(sample_data, ["colorList"])

2. 批量生成二进制标识列
   定义需要检查的颜色集合，遍历每个颜色生成对应列：

target_colors = ["Red", "Blue", "Green", "Black"]

for color in target_colors:
    # 判断列表是否包含颜色，将布尔结果转成整数（True→1，False→0）
    df = df.withColumn(color, array_contains(col("colorList"), color).cast("integer"))

3. 查看最终结果
   执行df.show(truncate=False)后，输出与预期一致：

+------------------------------+---+----+-----+-----+
|colorList                     |Red|Blue|Green|Black|
+------------------------------+---+----+-----+-----+
|[Red, Blue, Green, Black]     |1  |1   |1    |1    |
|[Blue, Green]                 |0  |1   |1    |0    |
+------------------------------+---+----+-----+-----+

说明 #

• array_contains(col, value)：专门用于判断数组类型列是否包含指定元素，返回布尔值
• .cast("integer")：直接将布尔值转为整数，写法比when/otherwise更简洁高效

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