如何在PySpark中不使用explode实现数组元素级减法操作

方法1：使用内置函数zip_with（Spark 3.0+推荐） #

Spark 3.0及以上版本提供的zip_with函数，能直接对两个数组的对应元素应用自定义逻辑，完全不需要explode，这是效率最高的实现方式，优先推荐。

示例代码：

from pyspark.sql import SparkSession
from pyspark.sql.functions import zip_with

# 初始化SparkSession
spark = SparkSession.builder.appName("ArraySubtraction").getOrCreate()

# 创建测试DataFrame
data = [( [2,4], [1,1] )]
df = spark.createDataFrame(data, ["col1", "col2"])

# 执行元素级减法
result_df = df.withColumn("subtracted", zip_with("col1", "col2", lambda x, y: x - y))

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

输出结果：

+------+------+----------+
|col1  |col2  |subtracted|
+------+------+----------+
|[2, 4]|[1, 1]|[1, 3]    |
+------+------+----------+

zip_with会将两个数组按位置配对，对每一对元素执行传入的lambda函数（此处为减法操作）。若两个数组长度不一致，默认以较短数组的长度为准，超出部分会被忽略；如果需要处理长度不一致的场景，可在lambda中添加判断逻辑。

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

如果你的Spark版本低于3.0，无法使用zip_with，可以自定义UDF来实现元素级减法。

示例代码：

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

# 初始化SparkSession
spark = SparkSession.builder.appName("ArraySubtractionUDF").getOrCreate()

# 创建测试DataFrame
data = [( [2,4], [1,1] )]
df = spark.createDataFrame(data, ["col1", "col2"])

# 定义元素级减法的UDF
def array_subtract(arr1, arr2):
    # 假设两个数组长度相同，若需处理长度不一致场景，可添加校验逻辑
    return [a - b for a, b in zip(arr1, arr2)]

# 注册UDF
array_subtract_udf = udf(array_subtract, ArrayType(IntegerType()))

# 应用UDF
result_df = df.withColumn("subtracted", array_subtract_udf("col1", "col2"))

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

输出结果与方法1一致。如果需要严格要求两个数组长度相同，可在UDF中添加校验：

def array_subtract(arr1, arr2):
    if len(arr1) != len(arr2):
        raise ValueError("Arrays must have the same length")
    return [a - b for a, b in zip(arr1, arr2)]

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