如何在PySpark窗口函数中基于动态条件排除行？

这个问题确实挺挠头的——常规的PySpark窗口函数根本没法直接处理这种动态依赖的行排除逻辑！毕竟窗口函数是基于固定的行范围来计算的，而你这里需要根据之前算出的feature结果，动态跳过那些标记为True的行，后续的中位数计算完全依赖前序的动态状态，这就超出了普通窗口函数的能力范围了。

先看看你当前代码的问题：你用了rowsBetween(-5, 0)来定义窗口，这是固定取当前行和前5行，不管这些行的feature是什么。但你需要的是，一旦某行的feature变成True，后续所有中位数计算都要跳过这行，普通窗口函数根本做不到这种动态过滤。

那该怎么解决呢？咱们可以用**递归CTE（公共表表达式）**来处理这种逐行依赖的逻辑，因为它能帮咱们维护一个动态的状态，记录哪些行是符合条件（feature=False）的，从而实现动态排除。

具体实现步骤 #

1. 先给数据加个连续序号

因为递归需要按顺序逐行处理，咱们得确保数据严格按id排序，并且有一个连续的序号列：

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

# 给数据添加连续的seq列，确保按id顺序处理
df = df.withColumn("seq", F.row_number().over(Window.orderBy("id")))

2. 自定义中位数UDF（符合你要的“取lower值”需求）

PySpark自带的percentile_approx对数组的处理可能没法直接满足你“平局取lower值”的要求，所以咱们写个简单的UDF：

from pyspark.sql.types import DoubleType

def median_lower(arr):
    if not arr:
        return None
    sorted_arr = sorted(arr)
    n = len(sorted_arr)
    # 取lower中位数：比如6个元素取第3个（0-based索引是2）
    mid_idx = (n - 1) // 2
    return sorted_arr[mid_idx]

median_lower_udf = F.udf(median_lower, DoubleType())
# 注册到Spark SQL里方便后续使用
spark.udf.register("median_lower", median_lower_udf)

3. 用递归CTE实现动态计算

递归CTE分为两部分：基例（处理第一行）和递归步骤（逐行处理后续数据，维护有效行的状态）：

# 开启递归CTE支持（Spark 3.0+可能需要，旧版本可能要设置这个参数）
spark.sql("SET spark.sql.legacy.recursiveCTE.enabled=true")

cte_query = """
WITH RECURSIVE cte AS (
    -- 基例：处理第一行，初始化状态
    SELECT 
        id,
        value,
        array(value) AS valid_values,
        CAST(value AS DOUBLE) AS median_value,
        CAST(value > 35 AS BOOLEAN) AS feature
    FROM df WHERE seq = 1
    
    UNION ALL
    
    -- 递归步骤：逐行处理，动态维护有效行列表
    SELECT 
        curr.id,
        curr.value,
        -- 如果当前行的feature为False，就把它加入有效列表，只保留最近5个（避免数组过大）
        CASE WHEN curr_median <= 35 THEN
            SLICE(array_prev.valid_values || curr.value, GREATEST(1, SIZE(array_prev.valid_values || curr.value) - 5), 5)
        ELSE
            -- feature为True，不加入有效列表，后续计算会跳过这行
            array_prev.valid_values
        END AS valid_values,
        curr_median AS median_value,
        CAST(curr_median > 35 AS BOOLEAN) AS feature
    FROM (
        SELECT 
            curr.*,
            array_prev.valid_values,
            -- 用自定义UDF计算基于有效列表+当前值的中位数
            median_lower(array_prev.valid_values || curr.value) AS curr_median
        FROM df curr
        JOIN cte array_prev ON curr.seq = array_prev.seq + 1
    ) t
)
-- 最终只取需要的列
SELECT id, value, median_value, feature FROM cte ORDER BY id
"""

# 执行查询得到结果
result_df = spark.sql(cte_query)
result_df.show()

这个方案的核心逻辑 #

• 咱们维护了一个valid_values数组，专门存储那些feature=False的行的value，并且只保留最近5个（因为你只需要前5个值来计算中位数，这样能避免数组无限增大）。
• 每一行的中位数计算，都是基于之前所有有效行的value加上当前行的value，然后判断feature：如果中位数超过35，就不把当前行加入有效数组，后续的中位数计算就会自动跳过这行。
• 递归CTE的优势就是能逐行传递状态，完美解决这种“前序结果影响后续计算”的动态逻辑。

为什么常规窗口函数不行？ #

常规窗口函数的窗口范围是静态定义的——要么是固定的行数（rowsBetween），要么是固定的数值范围（rangeBetween），它没法根据计算出来的动态列（比如你的feature）来实时调整窗口内包含哪些行。而递归CTE是逐行处理，每一步都能基于上一步的结果来调整后续的计算数据源，正好匹配你的需求。

备注：内容来源于stack exchange，提问作者user29963762