Python/PySpark如何生成数组列中相邻数值的差值数组

Python原生实现 #

你说的迭代逻辑直接对应列表推导式写法，也可以用更简洁的zip配对写法实现：

• 基础列表推导式实现

def calc_adjacent_diff(arr):
    # 长度小于2的数组直接返回空数组，可按需调整规则
    return [arr[i] - arr[i-1] for i in range(1, len(arr))] if len(arr) >=2 else []

# 测试
test_arr = [0,80,160,220]
print(calc_adjacent_diff(test_arr)) # 输出 [80, 80, 60]

• 简化zip写法

def calc_adjacent_diff(arr):
    return [b - a for a, b in zip(arr, arr[1:])] if len(arr) >=2 else []

如果是处理Pandas DataFrame中的数组列，直接用apply调用上面的函数即可：

import pandas as pd
df = pd.DataFrame({"num_arr": [[0,80,160,220], [10,20,35,50]]})
df["diff_arr"] = df["num_arr"].apply(calc_adjacent_diff)

PySpark实现 #

方案1：内置高阶函数实现（Spark 2.4+ 推荐，性能比UDF好） #

直接用transform+sequence操作数组下标完成计算：

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

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

# 构造测试数据
data = [([0,80,160,220],), ([10,20,35,50],)]
df = spark.createDataFrame(data, ["num_arr"])

# 计算相邻差值
df = df.withColumn("diff_arr", expr(
    "transform(sequence(1, size(num_arr)-1), i -> num_arr[i] - num_arr[i-1])"
))

df.show(truncate=False)

输出结果：

+------------------+------------+
|num_arr           |diff_arr    |
+------------------+------------+
|[0, 80, 160, 220] |[80, 80, 60]|
|[10, 20, 35, 50]  |[10, 15, 15]|
+------------------+------------+

方案2：UDF实现（兼容低版本Spark） #

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

# 定义UDF，数组元素是浮点型的话把返回类型改成DoubleType即可
diff_udf = udf(calc_adjacent_diff, ArrayType(IntegerType()))

df = df.withColumn("diff_arr", diff_udf("num_arr"))

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