如何在PySpark DataFrame中实现数据分区及单分区记录数限制

你提的这个需求完全可以通过PySpark实现，核心使用PySpark的窗口函数功能即可完成全部操作，具体实现步骤和代码如下：

核心实现逻辑 #

• 按city、state、cuisine_name三个字段做窗口分区
• 每个分区内按stars、review_count两个字段排序（示例默认高星级、高评论数优先，可自定义正序/倒序）
• 给分区内每条数据生成行号，筛选行号小于等于限制值的记录，即可实现每个分区返回固定数量的结果

完整代码示例 #

# 导入依赖包
from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, col

# 初始化SparkSession
spark = SparkSession.builder.appName("partition_topn").getOrCreate()

# 这里替换成你自己的数据集读取逻辑，比如读csv、parquet等，示例假设你已经加载数据到df变量
# df = spark.read.csv("your_data_path", header=True, inferSchema=True)

# 定义窗口规则
window_spec = Window.partitionBy("city", "state", "cuisine_name") \
    .orderBy(col("stars").desc(), col("review_count").desc())

# 给每个分区内的记录标记行号
df_with_row_num = df.withColumn("row_num", row_number().over(window_spec))

# 自定义每个分区最多返回的记录数，这里示例设为2，可按需调整
top_n = 2
# 筛选每个分区的topN记录，删除辅助行号列
result_df = df_with_row_num.filter(col("row_num") <= top_n).drop("row_num")

# 输出结果查看
result_df.show()

可选调整说明 #

• 排序规则调整：如果需要按星级升序、评论数升序排序，把orderBy里的.desc()改为.asc()即可
• 并列记录处理：如果要保留排序值相同的并列记录，把row_number()替换为rank()或dense_rank()即可：
  • row_number()：相同排序值的记录也会分配不同的连续行号，无重复
  • rank()：相同排序值的记录行号相同，会跳过后续行号（比如2条并列第1，下1条行号为3）
  • dense_rank()：相同排序值的记录行号相同，不跳过后续行号（比如2条并列第1，下1条行号为2）

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