PySpark基于数组列计算偏移量：transform方法报错求解决方案

需求 #

基于数组列Col1计算连续偏移区间，要求不使用UDF，尝试用transform函数实现时出现错误，需解决。

示例输入输出 #

错误代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, explode, expr

spark = SparkSession.builder \
.appName("Calculate Offset Column") \
.getOrCreate()

data = [([3.0, 4.0, 6.0, 2.0, 1.0],),
    ([10.0, 5.0, 4.0, 3.0, 2.0],)]

df = spark.createDataFrame(data, ["Col1"])
df = df.withColumn("Offsets", 
               expr("""transform(Col1, (x, i) -> struct(coalesce(sum(Col1) over (order by i rows between unbounded preceding and current row) - x, 0) as start, 
                                                         sum(Col1) over (order by i rows between unbounded preceding and current row) as end))"""))

报错信息 #

Error: 解析属性时发现问题，算子中的Window部分缺少属性i#462（来自Col1#454）。详细错误：!Window [Col1#454, transform(Col1#454, lambdafunction(struct(start, coalesce((sum(cast(Col1 as double)) windowspecdefinition(lambda i#462 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$())) - lambda x#461), cast(0 as double)), end, sum(cast(Col1 as double)) windowspecdefinition(lambda i#462 ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), currentrow$()))), lambda x#461, lambda i#462, false)) AS Offsets#458], [lambda i#462 ASC NULLS FIRST].;

错误原因 #

transform函数中的i是数组内部的索引变量，属于lambda作用域内的参数，而Window函数是作用于DataFrame行级别的操作，无法识别该内部变量，因此触发属性缺失错误。

可行解决方案 #

方案1：展开数组计算后聚合（直观易懂） #

通过展开数组、计算累计值、生成偏移区间后重新聚合，避免使用UDF：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, explode_outer, row_number, sum as sum_, count, struct, collect_list
from pyspark.sql.window import Window

spark = SparkSession.builder \
.appName("Calculate Offset Column") \
.getOrCreate()

data = [([3.0, 4.0, 6.0, 2.0, 1.0],),
    ([10.0, 5.0, 4.0, 3.0, 2.0],)]

df = spark.createDataFrame(data, ["Col1"])

# 添加行唯一标识，用于后续聚合
df = df.withColumn("row_id", row_number().over(Window.orderBy(col("Col1"))))

# 展开数组并添加元素索引
df_exploded = df.select(
    "row_id",
    col("Col1"),
    explode_outer(col("Col1")).alias("val"),
    row_number().over(Window.partitionBy("row_id").orderBy(col("Col1"))).alias("idx")
)

# 计算累计求和与累计计数
window_spec = Window.partitionBy("row_id").orderBy("idx").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df_cum = df_exploded.withColumn(
    "cum_sum", sum_("val").over(window_spec)
).withColumn(
    "cum_count", count("val").over(window_spec)
)

# 计算每个元素的偏移区间
df_offset = df_cum.withColumn(
    "offset",
    struct(
        (col("cum_sum") - col("val") + (col("cum_count") - 1)).alias("start"),
        (col("cum_sum") + (col("cum_count") - 1)).alias("end")
    )
)

# 聚合回数组格式
df_result = df_offset.groupBy("row_id", "Col1").agg(
    collect_list("offset").alias("Offsets")
).drop("row_id")

df_result.show(truncate=False)

方案2：纯高阶函数实现（无需展开数组） #

利用aggregate计算数组前缀累计值，结合transform生成偏移区间，全程使用Spark内置高阶函数：

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

spark = SparkSession.builder \
.appName("Calculate Offset Column") \
.getOrCreate()

data = [([3.0, 4.0, 6.0, 2.0, 1.0],),
    ([10.0, 5.0, 4.0, 3.0, 2.0],)]

df = spark.createDataFrame(data, ["Col1"])

# 用aggregate计算前缀累计和与索引，再通过transform生成偏移区间
df = df.withColumn(
    "Offsets",
    expr("""
        transform(
            Col1,
            (x, i) -> {
                val prefix = aggregate(
                    slice(Col1, 1, i+1),
                    cast(struct(0.0 as sum, 0 as idx) as struct<sum:double, idx:int>),
                    (acc, y) -> struct(acc.sum + y as sum, acc.idx + 1 as idx)
                );
                struct(
                    (prefix.sum - x + (prefix.idx - 1)) as start,
                    (prefix.sum + (prefix.idx - 1)) as end
                )
            }
        )
    """)
)

df.show(truncate=False)

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