PySpark中如何生成元素级布尔数组：判断col2元素是否存在于col1

要实现PySpark中检查数组列元素存在性并生成对应0/1数组的需求，不能直接用NumPy的函数（PySpark是分布式计算框架，无法直接调用本地NumPy方法），推荐使用Spark内置函数或自定义UDF两种方式：

方法1：使用Spark内置函数（推荐，性能更优） #

Spark 2.4及以上版本支持transform函数遍历数组，结合array_contains判断元素是否存在，代码示例如下：

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

# 初始化Spark会话
spark = SparkSession.builder.appName("ArrayElementCheck").getOrCreate()

# 创建示例数据
sample_data = [
    (["item1", "item2", "item3"], ["item5", "item2", "item3", "item17"]),
    (["item3", "item5", "item6", "item9"], ["item3", "item2", "item9", "item5", "item12"])
]
df = spark.createDataFrame(sample_data, schema=["col1", "col2"])

# 添加结果列
df = df.withColumn(
    "result_col",
    transform(
        col("col2"),
        lambda element: when(array_contains(col("col1"), element), 1).otherwise(0)
    )
)

# 查看结果
df.show(truncate=False)

运行后输出结果：

+-------------------------+-------------------------------+---------------+
|col1                     |col2                           |result_col     |
+-------------------------+-------------------------------+---------------+
|[item1, item2, item3]    |[item5, item2, item3, item17]  |[0, 1, 1, 0]   |
|[item3, item5, item6, item9]|[item3, item2, item9, item5, item12]|[1, 0, 1, 1, 0]|
+-------------------------+-------------------------------+---------------+

方法2：自定义UDF（兼容旧版本Spark） #

如果你的Spark版本低于2.4，可以用自定义UDF实现：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType

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

sample_data = [
    (["item1", "item2", "item3"], ["item5", "item2", "item3", "item17"]),
    (["item3", "item5", "item6", "item9"], ["item3", "item2", "item9", "item5", "item12"])
]
df = spark.createDataFrame(sample_data, schema=["col1", "col2"])

# 定义判断逻辑的函数
def check_element_existence(col1_array, col2_array):
    return [1 if item in col1_array else 0 for item in col2_array]

# 注册UDF
check_udf = udf(check_element_existence, ArrayType(IntegerType()))

# 添加结果列
df = df.withColumn("result_col", check_udf(col("col1"), col("col2")))

df.show(truncate=False)

注意事项 #

• 优先使用内置函数，UDF需要序列化数据到Executor节点执行，性能不如原生Spark函数。
• 确保Spark版本支持对应函数（transform函数从Spark 2.4开始提供）。

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