Spark中基于分组前一行值生成final列的问题求助

需要在Spark DataFrame中生成名为final的新列，规则如下：

• 分组依据：colA、colB、colC、colD
• 组内按colE排序，仅colE值变化
• 组内第一行的final值 = value * pred
• 组内后续行的final值 = 前一行的final值 * 当前行的pred

当前代码仅能生成组内前两行的final值，第三行及以后均为null，需排查错误原因并给出解决方法。

相关代码与数据 #

创建SparkSession与示例数据 #

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("example") \
    .getOrCreate()

# 示例数据
data = [
    ("A", "2003-03-01", 1, 11, 1, 10, 0.1),
    ("A", "2003-03-01", 1, 11, 2, 10, 0.2),
    ("A", "2003-03-01", 1, 11, 3, 10, 0.3),
    ("A", "2003-03-01", 1, 11, 4, 10, 0.1),
    ("A", "2003-03-01", 1, 11, 5, 10, 0.2),
]

# 创建DataFrame
df = spark.createDataFrame(data, ["colA", "colB", "colC", "colD", "colE", "value", "pred"])

# 期望输出数据
output = [
    ("A", "2003-03-01", 1, 11, 1, 10, 0.1, 1),
    ("A", "2003-03-01", 1, 11, 2, 10, 0.2, 0.2),
    ("A", "2003-03-01", 1, 11, 3, 10, 0.3, 0.06),
    ("A", "2003-03-01", 1, 11, 4, 10, 0.1, 0.006),
    ("A", "2003-03-01", 1, 11, 5, 10, 0.2, 0.0012),
]

output_df = spark.createDataFrame(output, ["colA", "colB", "colC", "colD", "colE", "value", "pred", "final"])

当前错误实现 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

window_spec = Window.partition('colA', 'colB', 'colC', 'colD').orderBy('colE')
# 生成组内第一行的值
a1 = df.withColumn('final', F.when(F.lag('colE').over(window_spec).isNull(), F.col('pred')*F.col('value')))
a1 = a1.withColumn('final', F.when(F.col('final').isNotNull(), F.col('final'))
                    .otherwise(F.lag(F.col('final')).over(window_spec) * F.col('pred')))

错误输出 #

+----+----------+----+----+----+-----+----+-----+
|colA|      colB|colC|colD|colE|value|pred|final|
+----+----------+----+----+----+-----+----+-----+
|   A|2003-03-01|   1|  11|   1|   10| 0.1|  1.0|
|   A|2003-03-01|   1|  11|   2|   10| 0.2|  0.2|
|   A|2003-03-01|   1|  11|   3|   10| 0.3| null|
|   A|2003-03-01|   1|  11|   4|   10| 0.1| null|
|   A|2003-03-01|   1|  11|   5|   10| 0.2| null|
+----+----------+----+----+----+-----+----+-----+

Spark的DataFrame是不可变数据结构，所有列操作都是基于原DataFrame的快照生成新列，无法在单次计算流程中迭代引用刚生成的列的实时值。

你的代码中：

1. 第一次withColumn生成的final列仅第一行有值，其余行均为null
2. 第二次withColumn中，lag('final')引用的是第一次生成的final列（仅第一行非空），而非更新后的final列
3. 第三行及以后的lag('final')取到的都是null，与pred相乘后仍为null，导致后续行结果异常

final列的本质是value乘以组内从第一行到当前行的pred累积乘积，可通过两种方式实现：

方法一：利用对数转累加计算累积乘积（适用于所有Spark版本） #

通过将乘积转换为对数的累加，再转指数还原，实现累积乘积计算：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

window_spec = Window.partitionBy('colA', 'colB', 'colC', 'colD').orderBy('colE')

df_final = df.withColumn(
    # 计算组内pred的累积乘积
    'cumulative_pred',
    F.exp(F.sum(F.log(F.col('pred'))).over(window_spec))
).withColumn(
    # 生成final列
    'final',
    F.col('value') * F.col('cumulative_pred')
).drop('cumulative_pred')

df_final.show()

注意：此方法要求pred必须为正数，若存在0或负数，需额外用when逻辑处理。

方法二：使用Spark 3.0+的Aggregate窗口函数 #

Spark 3.0及以上支持aggregate窗口函数，可直接迭代计算累积乘积：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

window_spec = Window.partitionBy('colA', 'colB', 'colC', 'colD').orderBy('colE')

df_final = df.withColumn(
    'final',
    # 聚合组内从第一行到当前行的pred，计算累积乘积
    F.aggregate(
        F.collect_list('pred').over(window_spec.rangeBetween(Window.unboundedPreceding, 0)),
        F.lit(1.0),  # 初始值
        lambda acc, x: acc * x  # 迭代逻辑：累积乘积
    ) * F.col('value')
)

df_final.show()

两种方法均可得到符合预期的final列结果。

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