Spark SQL 3.3版本LEAD窗口函数是否支持表达式作为偏移量？

问题描述 #

在Spark 3.3中运行包含LEAD函数的查询时，触发如下错误：

AnalysisException: cannot resolve 'lead(y.dense_rnk1, ((y.last_val1 - y.first_val1) + 1), NULL)' due to data type mismatch: Offset expression '((last_val1#18047 - first_val1#18046) + 1)' must be a literal.

虽然Spark SQL文档说明2.3版本后允许将表达式作为LEAD/LAG的偏移量，但实际运行时仍报错。已尝试将偏移量表达式中的字段转为整数、使用COALESCE处理NULL值，问题未解决。

原查询核心代码：

SELECT
  y.*,
  LEAD(y.dense_rnk1, (y.last_val1 - y.first_val1 + 1), NULL) OVER (
    PARTITION BY y.Shipment_Tracking_Nr,
    y.Tracking_Item_Unique_Id
    ORDER BY
      y.Reporting_Dt
  ) AS lead_dense_rnk1
FROM
  (
    SELECT
      x.*,
      FIRST_VALUE(x.row_num1) OVER(
        PARTITION BY x.Shipment_Tracking_Nr,
        x.Tracking_Item_Unique_Id,
        x.dense_rnk1
        ORDER BY
          x.row_num1 RANGE BETWEEN UNBOUNDED PRECEDING
          AND CURRENT ROW
      ) AS first_val1,
      LAST_VALUE(x.row_num1) OVER(
        PARTITION BY x.Shipment_Tracking_Nr,
        x.Tracking_Item_Unique_Id,
        x.dense_rnk1
        ORDER BY
          x.row_num1 RANGE BETWEEN UNBOUNDED PRECEDING
          AND UNBOUNDED FOLLOWING
      ) AS last_val1
    FROM x
  ) y

问题原因 #

Spark的LEAD/LAG函数对表达式偏移量的支持有隐性限制：偏移量表达式必须在当前窗口分区内是全局常量。也就是说，同一个窗口分区内的所有行，偏移量的值必须完全相同。

在当前查询中，last_val1 - first_val1 +1是基于dense_rnk1子分区计算的常量，但外层LEAD函数的窗口分区是Shipment_Tracking_Nr, Tracking_Item_Unique_Id——同一个外层分区内可能包含多个dense_rnk1分组，导致偏移量在同一个LEAD窗口分区内不是固定值。Spark在分析期无法确认该表达式是窗口内常量，因此判定它不符合要求，抛出错误。

解决方案 #

方案1：调整LEAD窗口的分区规则 #

如果业务逻辑允许，将dense_rnk1加入LEAD函数的PARTITION BY列表，确保偏移量在新的窗口分区内是常量：

SELECT
  y.*,
  LEAD(y.dense_rnk1, (y.last_val1 - y.first_val1 + 1), NULL) OVER (
    PARTITION BY y.Shipment_Tracking_Nr,
    y.Tracking_Item_Unique_Id,
    y.dense_rnk1  -- 新增分区字段
    ORDER BY
      y.Reporting_Dt
  ) AS lead_dense_rnk1
FROM
  (
    SELECT
      x.*,
      FIRST_VALUE(x.row_num1) OVER(
        PARTITION BY x.Shipment_Tracking_Nr,
        x.Tracking_Item_Unique_Id,
        x.dense_rnk1
        ORDER BY
          x.row_num1 RANGE BETWEEN UNBOUNDED PRECEDING
          AND CURRENT ROW
      ) AS first_val1,
      LAST_VALUE(x.row_num1) OVER(
        PARTITION BY x.Shipment_Tracking_Nr,
        x.Tracking_Item_Unique_Id,
        x.dense_rnk1
        ORDER BY
          x.row_num1 RANGE BETWEEN UNBOUNDED PRECEDING
          AND UNBOUNDED FOLLOWING
      ) AS last_val1
    FROM x
  ) y

方案2：用自连接替代LEAD函数 #

如果不能调整窗口分区，可通过给数据添加行号，再自连接的方式实现动态偏移量的取值逻辑：

1. 先给外层查询的行添加全局行号，并计算偏移量：

WITH ranked_data AS (
  SELECT
    y.*,
    ROW_NUMBER() OVER (
      PARTITION BY y.Shipment_Tracking_Nr, y.Tracking_Item_Unique_Id
      ORDER BY y.Reporting_Dt
    ) AS row_idx,
    (y.last_val1 - y.first_val1 + 1) AS offset_val
  FROM
    (
      SELECT
        x.*,
        FIRST_VALUE(x.row_num1) OVER(
          PARTITION BY x.Shipment_Tracking_Nr,
          x.Tracking_Item_Unique_Id,
          x.dense_rnk1
          ORDER BY
            x.row_num1 RANGE BETWEEN UNBOUNDED PRECEDING
            AND CURRENT ROW
        ) AS first_val1,
        LAST_VALUE(x.row_num1) OVER(
          PARTITION BY x.Shipment_Tracking_Nr,
          x.Tracking_Item_Unique_Id,
          x.dense_rnk1
          ORDER BY
            x.row_num1 RANGE BETWEEN UNBOUNDED PRECEDING
            AND UNBOUNDED FOLLOWING
        ) AS last_val1
      FROM x
    ) y
)

2. 自连接获取目标偏移量对应的dense_rnk1值：

SELECT
  main.*,
  lead_data.dense_rnk1 AS lead_dense_rnk1
FROM ranked_data main
LEFT JOIN ranked_data lead_data
  ON main.Shipment_Tracking_Nr = lead_data.Shipment_Tracking_Nr
  AND main.Tracking_Item_Unique_Id = lead_data.Tracking_Item_Unique_Id
  AND lead_data.row_idx = main.row_idx + main.offset_val

补充说明 #

• 检查first_val1和last_val1的计算逻辑：你已指定RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING给LAST_VALUE，确保能获取到整个dense_rnk1分组的最后值，这部分是正确的。
• 如果偏移量可能为负数或超出窗口范围，自连接时会返回NULL，和LEAD函数的默认行为一致。

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