请求将指定SAS代码转换为Databricks兼容的PySpark SQL/Python代码

原SAS代码逻辑解析 #

原SAS代码通过数组和lag函数实现12行滑动窗口求和：对15组触发列（PX_fl_*系列），每组对应一个目标列，目标列值为当前行触发值 + 前11行触发值的总和。本质是标记当前行及过去11行内是否有触发事件（值≥1即表示在12个月周期内有事件）。

注意：Spark中必须确保数据按用户/主体列分区、时间列排序，否则滑动窗口的结果会不符合预期（SAS默认按输入顺序处理，Spark需要显式指定分区和排序规则）。

方式1：PySpark DataFrame实现 #

假设你的数据有：

• 分区列：patient_id（按主体分组）
• 时间排序列：month_dt（按月份排序，确保时间顺序）
• 触发列列表与目标列列表一一对应（修正原SAS中目标列重复的笔误，确保15组一一匹配）

from pyspark.sql.window import Window
from pyspark.sql.functions import sum, col, when

# 读取源表
df = spark.table("out2204.pxrx_tot_v2")

# 定义触发列和目标列（一一对应，修正原SAS中目标列重复问题）
trigger_cols = [
    "PX_fl_1", "PX_fl_2", "PX_fl_3", "PX_fl_4", "PX_fl_5",
    "PX_fl_6", "PX_fl_7", "PX_fl_8", "PX_fl_9", "PX_fl_10",
    "PX_fl_11", "PX_fl_12", "PX_fl_13", "PX_fl_TPN", "PX_fl_HYD"
]
target_cols = [
    "PN", "TPN_1", "TPN_2", "TPN_3", "TPN_4",
    "TPN_5", "TPN_6", "TPN_7", "TPN_8", "TPN_9",
    "TPN_10", "TPN_11", "TPN_T", "PN_TPN", "HYD"
]

# 定义滑动窗口：按患者分区，按月份排序，窗口范围为当前行+前11行
window_spec = Window.partitionBy("patient_id").orderBy("month_dt").rowsBetween(-11, 0)

# 对每组列计算窗口求和，生成目标列
for trigger_col, target_col in zip(trigger_cols, target_cols):
    # 计算窗口内的和，可选：转成1/0标记（只要有事件就为1）
    df = df.withColumn(
        target_col,
        when(sum(col(trigger_col)).over(window_spec) >= 1, 1).otherwise(0)
        # 若需要保留求和值（叠加效果），直接用sum(col(trigger_col)).over(window_spec)
    )

# 保存结果表
df.write.mode("overwrite").saveAsTable("pxrx_tot2")

方式2：Databricks SQL实现 #

同样基于滑动窗口函数，语法更简洁：

CREATE OR REPLACE TABLE pxrx_tot2 AS
SELECT
    *,
    -- 第1组：PX_fl_1 → PN
    CASE WHEN SUM(PX_fl_1) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS PN,
    -- 第2组：PX_fl_2 → TPN_1
    CASE WHEN SUM(PX_fl_2) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_1,
    -- 依次添加剩余13组列，保持触发列与目标列一一对应
    CASE WHEN SUM(PX_fl_3) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_2,
    CASE WHEN SUM(PX_fl_4) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_3,
    CASE WHEN SUM(PX_fl_5) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_4,
    CASE WHEN SUM(PX_fl_6) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_5,
    CASE WHEN SUM(PX_fl_7) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_6,
    CASE WHEN SUM(PX_fl_8) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_7,
    CASE WHEN SUM(PX_fl_9) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_8,
    CASE WHEN SUM(PX_fl_10) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_9,
    CASE WHEN SUM(PX_fl_11) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_10,
    CASE WHEN SUM(PX_fl_12) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_11,
    CASE WHEN SUM(PX_fl_13) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS TPN_T,
    CASE WHEN SUM(PX_fl_TPN) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS PN_TPN,
    CASE WHEN SUM(PX_fl_HYD) OVER (PARTITION BY patient_id ORDER BY month_dt ROWS BETWEEN 11 PRECEDING AND CURRENT ROW) >= 1 THEN 1 ELSE 0 END AS HYD
FROM out2204.pxrx_tot_v2;

关键注意事项 #

1. 分区与排序：必须指定partitionBy（如患者ID）和orderBy（如时间列），否则窗口函数会在全表范围内计算，结果错误。
2. 列对应关系：原SAS代码中目标列数组存在重复（两个PN），需根据实际业务修正为一一对应的15组列。
3. 叠加效果：若需要保留求和的叠加值（而非仅1/0标记），直接去掉CASE WHEN，保留SUM(...) OVER (...)即可。

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