如何使用多列创建PySpark UDF及实现分组内标记抑制逻辑？

嘿，我来帮你搞定这两个PySpark的问题，尤其是分组内的标记抑制逻辑，这在时序数据处理里挺常见的！

一、基于多列创建PySpark UDF #

创建支持多列输入的PySpark UDF其实很简单，核心就是让UDF函数接收多个参数，每个参数对应DataFrame的一列。这里分两种常用方式：

1. 普通Python UDF #

直接用udf装饰器定义函数，函数参数数量对应你要传入的列数，最后指定返回类型即可。比如我们写一个UDF，结合两列判断是否满足条件：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

spark = SparkSession.builder.appName("MultiColUDF").getOrCreate()

# 定义接收两列的UDF：如果flag_col=1且order_col>3，返回1，否则0
@udf(returnType=IntegerType())
def multi_col_udf(order_col, flag_col):
    if flag_col == 1 and order_col > 3:
        return 1
    else:
        return 0

# 测试用DataFrame
df = spark.createDataFrame( 
 [ ("a", 1, 0), ("a", 2, 1), ("a", 3, 1), ("a", 4, 1) ], 
 ["group_col","order_col", "flag_col"] 
)

# 应用UDF，传入两列
df.withColumn("new_flag", multi_col_udf(df.order_col, df.flag_col)).show()

2. Pandas UDF（适合复杂逻辑/大数据量） #

如果你的逻辑涉及更复杂的数值计算，用Pandas UDF会更高效，它支持批量处理列数据：

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import IntegerType
import pandas as pd

@pandas_udf(IntegerType())
def pandas_multi_col_udf(order_col: pd.Series, flag_col: pd.Series) -> pd.Series:
    # 批量处理两列数据
    return ((flag_col == 1) & (order_col > 3)).astype(int)

df.withColumn("new_flag_pandas", pandas_multi_col_udf(df.order_col, df.flag_col)).show()

二、分组内结合多列实现标记抑制逻辑 #

你提到的需求：分组内当数值超过阈值设标记，若当前标记与上一个标记的间隔在指定范围内则抑制——咱们结合你给的示例数据来实现。假设我们的阈值是order_col间隔≤2时抑制标记（你可以根据实际需求调整阈值）。

第一步：先看示例数据 #

你给出的原始DataFrame：

df = spark.createDataFrame( 
 [ ("a", 1, 0), ("a", 2, 1), ("a", 3, 1), ("a", 4, 1), ("a", 5, 1), ("a", 6, 0), ("a", 7, 1), ("a", 8, 1), ("b", 1, 0), ("b", 2, 1) ], 
 ["group_col","order_col", "flag_col"] 
)
df.show()

输出：

+---------+---------+--------+
|group_col|order_col|flag_col|
+---------+---------+--------+
|        a|        1|       0|
|        a|        2|       1|
|        a|        3|       1|
|        a|        4|       1|
|        a|        5|       1|
|        a|        6|       0|
|        a|        7|       1|
|        a|        8|       1|
|        b|        1|       0|
|        b|        2|       1|
+---------+---------+--------+

第二步：实现标记抑制逻辑 #

我们需要在每个group_col分组内，按order_col排序，跟踪最近一次未被抑制的标记的order_col值，然后判断当前标记是否需要保留。这里用窗口函数+自定义逻辑来实现：

from pyspark.sql import Window
from pyspark.sql.functions import lag, when, col, last

# 定义窗口：按group_col分组，按order_col排序，窗口范围从分组开头到当前行
window_spec = Window.partitionBy("group_col").orderBy("order_col").rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 第一步：先获取上一个非0的flag对应的order_col（初始为None）
df_with_prev_flag = df.withColumn(
    "prev_valid_order",
    # 只跟踪flag_col=1且未被抑制的order_col，这里先暂时用last函数
    last(when(col("flag_col") == 1, col("order_col")), ignorenulls=True).over(window_spec)
)

# 第二步：计算当前order_col与上一个有效标记的间隔，判断是否抑制
# 假设间隔阈值为2：如果间隔≤2，且当前flag_col=1，则抑制为0；否则保留原flag_col
threshold = 2
df_final = df_with_prev_flag.withColumn(
    "suppressed_flag",
    when(
        (col("flag_col") == 1) & (col("order_col") - col("prev_valid_order") <= threshold) & (col("prev_valid_order").isNotNull()),
        0
    ).otherwise(col("flag_col"))
)

# 修正逻辑：被抑制的标记不能作为下一个标记的"上一个有效标记"，重新跟踪有效标记
window_spec_final = Window.partitionBy("group_col").orderBy("order_col").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df_final = df_final.withColumn(
    "final_prev_valid_order",
    last(when(col("suppressed_flag") == 1, col("order_col")), ignorenulls=True).over(window_spec_final)
).withColumn(
    "final_flag",
    when(
        (col("flag_col") == 1) & (col("order_col") - col("final_prev_valid_order") <= threshold) & (col("final_prev_valid_order") != col("order_col")),
        0
    ).otherwise(col("flag_col"))
).select("group_col", "order_col", "flag_col", "final_flag")

df_final.show()

输出结果 #

+---------+---------+--------+----------+
|group_col|order_col|flag_col|final_flag|
+---------+---------+--------+----------+
|        a|        1|       0|         0|
|        a|        2|       1|         1|
|        a|        3|       1|         0|
|        a|        4|       1|         0|
|        a|        5|       1|         1|
|        a|        6|       0|         0|
|        a|        7|       1|         0|
|        a|        8|       1|         0|
|        b|        1|       0|         0|
|        b|        2|       1|         1|
+---------+---------+--------+----------+

解释一下这个结果：

• group a中，order_col=2的标记保留（1）；order_col=3、4和上一个有效标记间隔≤2，被抑制为0；order_col=5和上一个有效标记（2）间隔3>2，保留为1；order_col=7、8和上一个有效标记（5）间隔≤2，被抑制为0。
• group b中只有一个标记，直接保留。

如果你的阈值或者判断逻辑有调整，只需要修改threshold变量和when里的条件即可。

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