求PySpark中高效结合排序过滤，按列值取指定行数的实现方案

需求明确 #

从DataFrame中筛选出指定列值最小的num_of_rows行：

• 若总行数≤指定数量，直接返回全部数据
• 保留筛选出的行在原DataFrame中的原始顺序

基础高效实现 #

适用于大部分场景，兼顾准确性和效率：

from pyspark.sql.functions import monotonically_increasing_id, col, row_number
from pyspark.sql.window import Window

def sort_and_filter_based_on_column(df, column, num_of_rows):
    total_rows = df.count()
    if total_rows <= num_of_rows:
        return df
    
    # 新增自增列记录原始行顺序
    df_with_order = df.withColumn("original_idx", monotonically_increasing_id())
    
    # 用窗口函数标记值最小的num_of_rows行
    rank_window = Window.orderBy(col(column).asc())
    df_ranked = df_with_order.withColumn("rank", row_number().over(rank_window))
    
    # 筛选目标行并恢复原始顺序
    result = (
        df_ranked.filter(col("rank") <= num_of_rows)
        .orderBy("original_idx")
        .drop("original_idx", "rank")
    )
    
    return result

超大规模数据优化方案 #

如果处理TB级别的数据集，全局排序的窗口函数开销过高，可以用分位数阈值减少计算量：

from pyspark.sql.functions import monotonically_increasing_id, col

def sort_and_filter_based_on_column(df, column, num_of_rows):
    total_rows = df.count()
    if total_rows <= num_of_rows:
        return df
    
    # 计算第num_of_rows小值的精确阈值（approxQuantile第三个参数设为0.0保证精确）
    quantile = num_of_rows / total_rows
    threshold = df.approxQuantile(column, [quantile], 0.0)[0]
    
    # 先过滤出值≤阈值的行
    filtered_df = df.filter(col(column) <= threshold)
    filtered_count = filtered_df.count()
    
    # 若过滤后行数超标，再小范围排序取数并恢复原顺序
    if filtered_count > num_of_rows:
        df_with_order = filtered_df.withColumn("original_idx", monotonically_increasing_id())
        result = (
            df_with_order.orderBy(col(column).asc())
            .limit(num_of_rows)
            .orderBy("original_idx")
            .drop("original_idx")
        )
    else:
        result = filtered_df
    
    return result

验证示例 #

先初始化测试DataFrame：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("TestSortFilter").getOrCreate()
test_data = [("a",4), ("b",5), ("c",3), ("d",1), ("e",2)]
df = spark.createDataFrame(test_data, ["keys", "values"])

调用函数验证：

# 取values最小的3行
sort_and_filter_based_on_column(df, "values", 3).show()
# 输出符合预期：
# +----+------+
# |keys|values|
# +----+------+
# |   c|     3|
# |   d|     1|
# |   e|     2|
# +----+------+

# 取values最小的2行
sort_and_filter_based_on_column(df, "values", 2).show()
# 输出：
# +----+------+
# |keys|values|
# +----+------+
# |   d|     1|
# |   e|     2|
# +----+------+

# 取全部5行
sort_and_filter_based_on_column(df, "values", 5).show()
# 输出原始全量数据：
# +----+------+
# |keys|values|
# +----+------+
# |   a|     4|
# |   b|     5|
# |   c|     3|
# |   d|     1|
# |   e|     2|
# +----+------+

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