如何在PySpark DataFrame中编写自动数据标注的Python函数

实现思路 #

• 针对3分类标注需求，通过分箱算法自动计算切分点替代硬编码阈值，支持等宽、等频两种分箱逻辑可选
• 等宽分箱：按字段取值的最大最小值平均切分成3段，适合取值分布均匀的场景
• 等频分箱：按数据分布的分位点切分，保证每个分类的样本量接近，适合分布倾斜的场景

完整实现代码 #

from pyspark.sql.functions import when, col, lit

def auto_label_delay(df, input_col="ARRIVAL_DELAY", output_col="multiClassArrDelay", bin_method="quantile"):
    """
    自动对到达延误字段做3分类标注
    :param df: 输入PySpark DataFrame
    :param input_col: 待标注的延误字段名
    :param output_col: 输出的标注列名
    :param bin_method: 分箱方式，quantile为等频分箱，width为等宽分箱
    :return: 新增标注列后的DataFrame
    """
    # 计算切分点
    if bin_method == "quantile":
        # 等频分箱：取33%、66%分位点作为切分边界
        splits = df.approxQuantile(input_col, [0.33, 0.66], 0.01)
        lower_split, upper_split = splits[0], splits[1]
    elif bin_method == "width":
        # 等宽分箱：按取值范围平均切3段
        stats = df.selectExpr(f"min({input_col}) as min_val", f"max({input_col}) as max_val").first()
        min_val, max_val = stats["min_val"], stats["max_val"]
        bin_width = (max_val - min_val) / 3
        lower_split = min_val + bin_width
        upper_split = min_val + 2 * bin_width
    else:
        raise ValueError("仅支持quantile或width两种分箱方式")
    
    # 生成标注列
    df = df.withColumn(output_col,
        when(col(input_col) < lower_split, lit(0))  # 早到
        .when((col(input_col) >= lower_split) & (col(input_col) <= upper_split), lit(1)) # 准点
        .otherwise(lit(2)) # 晚点
    )
    return df

# 调用示例
flightsDf = auto_label_delay(flightsDf, bin_method="quantile")

补充说明 #

• 等频分箱的第三个参数0.01是近似计算的误差容忍度，值越小精度越高、计算耗时越长，可根据实际数据量调整
• 如果需要自定义切分逻辑，直接修改splits对应的计算规则即可
• 原硬编码的固定阈值逻辑也可以封装到函数里做兜底，按需切换即可

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