PySpark如何实现列方向Lead/Lag完成相邻列滑动求和

列维度的滑动窗口聚合不需要做行列转换带来的额外shuffle，直接基于数量列的命名规则动态生成计算表达式即可，桶大小、滑动方向可灵活配置，性能远高于unpivot+行窗口+pivot的实现路径。以下分别给出PySpark和Spark SQL的可直接运行的实现，以桶大小=2的场景为例：Bucket_i = Quantity_i + Quantity_{i+1}。

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 初始化Spark会话（生产环境直接使用已有的SparkSession即可）
spark = SparkSession.builder.appName("column_sliding_calc").getOrCreate()

# --------------------------
# 1. 测试数据构造（可替换为自己的实际数据）
# --------------------------
sample_data = [
    (1, 10, 20, 30, 40, 50, 60),
    (2, 5, 15, 25, 35, 45, 55)
]
df = spark.createDataFrame(
    sample_data,
    schema=["id", "Quantity1", "Quantity2", "Quantity3", "Quantity4", "Quantity5", "Quantity6"]
)

# --------------------------
# 2. 参数配置
# --------------------------
BUCKET_SIZE = 2  # 可按需修改为5等其他桶大小
QTY_COL_PREFIX = "Quantity"
BUCKET_COL_PREFIX = "Bucket"

# --------------------------
# 3. 提取并排序所有数量列
# --------------------------
# 按列名末尾的数字升序排列，避免列顺序错乱
qty_columns = sorted(
    [col for col in df.columns if col.startswith(QTY_COL_PREFIX)],
    key=lambda x: int(x.replace(QTY_COL_PREFIX, ""))
)

# --------------------------
# 4. 动态生成所有Bucket列的计算逻辑
# --------------------------
bucket_calc_exprs = []
for idx in range(len(qty_columns) - BUCKET_SIZE + 1):
    # 取当前滑动窗口内的列求和，如需改成向前取（Lag效果），调整切片范围即可
    window_cols = qty_columns[idx : idx + BUCKET_SIZE]
    bucket_name = f"{BUCKET_COL_PREFIX}{idx+1}"
    # 如需处理空值，可改为 sum(F.coalesce(F.col(c), F.lit(0)) for c in window_cols)
    bucket_calc_exprs.append(
        sum(F.col(c) for c in window_cols).alias(bucket_name)
    )

# --------------------------
# 5. 生成结果
# --------------------------
result_df = df.select("*", *bucket_calc_exprs)
result_df.show()

运行后输出结果如下：

+---+---------+---------+---------+---------+---------+---------+-------+-------+-------+-------+-------+
| id|Quantity1|Quantity2|Quantity3|Quantity4|Quantity5|Quantity6|Bucket1|Bucket2|Bucket3|Bucket4|Bucket5|
+---+---------+---------+---------+---------+---------+---------+-------+-------+-------+-------+-------+
|  1|       10|       20|       30|       40|       50|       60|     30|     50|     70|     90|    110|
|  2|        5|       15|       25|       35|       45|       55|     20|     40|     60|     80|    100|
+---+---------+---------+---------+---------+---------+---------+-------+-------+-------+-------+-------+

逻辑和PySpark原生实现完全一致，动态拼接SQL语句即可：

# 提取排序后的数量列，逻辑同上文
qty_columns = sorted(
    [col for col in df.columns if col.startswith(QTY_COL_PREFIX)],
    key=lambda x: int(x.replace(QTY_COL_PREFIX, ""))
)

# 拼接Bucket字段的SELECT片段
bucket_select_clause = []
for idx in range(len(qty_columns) - BUCKET_SIZE + 1):
    window_cols = qty_columns[idx : idx + BUCKET_SIZE]
    sum_expr = " + ".join(window_cols)
    bucket_select_clause.append(f"{sum_expr} AS {BUCKET_COL_PREFIX}{idx+1}")

# 组装完整SQL
calc_sql = f"""
SELECT
    *,
    {",".join(bucket_select_clause)}
FROM source_table
"""

# 注册临时视图执行计算
df.createOrReplaceTempView("source_table")
sql_result = spark.sql(calc_sql)
sql_result.show()

• 如果需要实现最初需求的「汇总之前5个同类型字段」的Lag类效果，只需要把列切片的逻辑改为window_cols = qty_columns[max(0, idx - BUCKET_SIZE + 1) : idx + 1]，把BUCKET_SIZE设为5即可。
• 如果需要实现Lead类效果（取当前列之后N个字段聚合），调整切片起始位置即可，不需要修改核心逻辑。
• 该实现无额外shuffle开销，列数越多性能优势越明显。

注意：如果数量列存在空值，一定要在求和逻辑中加coalesce把空值转为0，否则任意一列空值会导致整个Bucket的求和结果为空。

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