PySpark如何从聚类预测结果的各簇中抽取n个样本并输出为DataFrame

全程基于Spark原生算子开发，无需转换为Pandas DataFrame，支持自定义簇筛选规则、单簇抽样逻辑，适配大数据量场景。

前置依赖导入 #

from pyspark.sql import Window, functions as F

通用函数定义 #

入参说明 #

• df: 带prediction簇标签列的输入Spark DataFrame
• top_n_clusters: 按簇实例数升序选取的目标簇数量
• samples_per_cluster: 每个簇需要返回的样本行数
• sort_exprs: 自定义簇内排序规则，传入pyspark.sql.Column组成的列表，和Spark原生orderBy()入参规则完全兼容

函数代码 #

def sample_clusters_by_rule(df, top_n_clusters, samples_per_cluster, sort_exprs):
    # 1. 统计各簇样本量，筛选符合要求的top n个小簇
    cluster_count_df = df.groupBy("prediction").agg(F.count("*").alias("cluster_size"))
    # 小数据量簇场景直接collect目标簇id，过滤非目标数据
    target_clusters = [row["prediction"] for row in cluster_count_df.orderBy(F.col("cluster_size").asc()).limit(top_n_clusters).collect()]
    filtered_df = df.filter(F.col("prediction").isin(target_clusters))
    
    # 2. 按簇分组，按自定义规则排序取前n个样本
    window_spec = Window.partitionBy("prediction").orderBy(*sort_exprs)
    result_df = filtered_df.withColumn("rank", F.row_number().over(window_spec)) \
                           .filter(F.col("rank") <= samples_per_cluster) \
                           .drop("rank")
    
    return result_df

若你的场景中簇的量级超过1万，可将目标簇的筛选逻辑改为join实现，避免driver侧collect数据过多：

# 替换原来的target_clusters collect+isin的逻辑即可
target_cluster_df = cluster_count_df.orderBy(F.col("cluster_size").asc()).limit(top_n_clusters).select("prediction")
filtered_df = df.join(target_cluster_df, on="prediction", how="inner")

使用示例 #

比如需要按簇大小升序取最小的5个簇，每个簇取3个x最小、y最大、z最大的样本，调用方式如下：

# 自定义排序规则，支持任意列的升降序组合
custom_sort_rule = [
    F.col("x").asc(),
    F.col("y").desc(),
    F.col("z").desc()
]

# 调用函数
result_df = sample_clusters_by_rule(
    df=your_input_df, # 替换为你自己的带prediction列的DataFrame
    top_n_clusters=5,
    samples_per_cluster=3,
    sort_exprs=custom_sort_rule
)

# 查看结果
result_df.show()

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