PySpark DataFrame按指定条件修剪失效，求代码排查与解决方法

核心问题分析 #

你的代码存在几个关键错误，导致无法正常过滤：

1. check_st函数逻辑错误：将Row对象转为set(i)得到的是列名集合，而非你需要匹配的字段值集合，因此aset.issubset(set(i))永远不会触发匹配，这就是没有打印There is a match的原因。
2. 误用collect()遍历DataFrame：collect()会把全量数据拉取到Driver节点，完全违背Spark分布式计算的设计，大数据场景下会直接崩溃，且效率极低。
3. rules_check规则判断错误：any((rules_1,rules_2))是判断两个列表是否非空，而非检查列表内所有条件是否同时满足，应该用all()来确保规则的四个条件全部成立。
4. 过滤逻辑不符合需求：all((check_st(...), rules_check(...)))要求两个函数都返回True才移除行，可能和你实际要过滤的规则不符。

推荐解决方案（Spark分布式实现） #

基于Spark窗口函数实现，完全利用分布式计算能力，适合所有数据规模：

from pyspark.sql import Window
from pyspark.sql.functions import col, datediff, max, min, count, when

def prune_df(sdf):
    # 按shop和email分区，确保同一用户同一店铺的记录放在一组
    window_spec = Window.partitionBy("shop", "email")
    
    # 计算分区内的时间范围，判断是否存在1天内的记录
    sdf_with_window = sdf.withColumn("min_dt", min("dt").over(window_spec)) \
                         .withColumn("max_dt", max("dt").over(window_spec)) \
                         .withColumn("days_diff", datediff(col("max_dt"), col("min_dt")))
    
    # 标记符合过滤规则的行：
    # 规则1：同组内至少2条declined记录，且时间差≤1天
    # 规则2：同组内同时存在accepted和declined记录，且时间差≤1天
    sdf_with_flags = sdf_with_window.withColumn(
        "has_declined_pair",
        when(
            (col("days_diff") <= 1) & 
            (count(when(col("decision") == "declined", 1)).over(window_spec) >= 2),
            True
        ).otherwise(False)
    ).withColumn(
        "has_accept_decline_pair",
        when(
            (col("days_diff") <= 1) & 
            (count(when(col("decision") == "accepted", 1)).over(window_spec) >= 1) & 
            (count(when(col("decision") == "declined", 1)).over(window_spec) >= 1),
            True
        ).otherwise(False)
    )
    
    # 过滤掉符合任一规则的行，清理临时字段
    pruned_sdf = sdf_with_flags.filter(
        ~col("has_declined_pair") & ~col("has_accept_decline_pair")
    ).drop("min_dt", "max_dt", "days_diff", "has_declined_pair", "has_accept_decline_pair")
    
    return pruned_sdf

原代码修正（仅适合小数据集） #

如果一定要用遍历方式（不推荐），需修正以下几点：

1. 修正check_st函数 #

def check_st(alist, aset):
    for i in alist:
        # 提取对应字段的值集合，而非列名
        i_values = set(list(map(i.__getitem__, [1,2,6,10])))  # 对应原代码的idx[1:]
        if aset.issubset(i_values):
            print(f'There is a match')
            return True
    return False

2. 修正rules_check函数 #

def rules_check(arow, alist):
    for i in alist:
        # 用all()确保所有规则条件同时满足，加abs避免时间差为负
        rules_1 = all([
            i['shop'] == arow['shop'],
            i['email'] == arow['email'],
            abs((i['dt'] - arow['dt']).days) <= 1,
            i['decision'] == 'declined' and arow['decision'] == 'declined'
        ])
        rules_2 = all([
            i['shop'] == arow['shop'],
            i['email'] == arow['email'],
            abs((i['dt'] - arow['dt']).days) <= 1,
            i['decision'] == 'accepted' and arow['decision'] == 'declined'
        ])
        if rules_1 or rules_2:
            return True
    return False

3. 修正prune_df函数 #

def prune_df(sdf):
    rem_idx = []
    store = []
    idx = [0,1,2,6,10]
    # 按dt排序，确保时间顺序正确
    for i in sdf.orderBy("dt").collect():
        current_set = set(list(map(i.__getitem__, idx))[1:])
        # 满足任一条件则标记移除
        if check_st(store, current_set) or rules_check(i, store):
            rem_idx.append(i['id'])
        else:
            store.append(i)
    pruned_sdf = sdf.filter(~sdf.id.isin(rem_idx))
    return pruned_sdf

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