PySpark如何用DataFrame API实现dense_rank窗口函数替代SQL查询

实现步骤 #

1. 导入依赖模块 #

首先需要导入窗口类和对应的窗口函数：

from pyspark.sql import Window
from pyspark.sql.functions import dense_rank

2. 定义窗口规范 #

对应SQL中OVER (PARTITION BY c1 ORDER BY value ASC)的逻辑，定义窗口规则：

window_spec = Window.partitionBy("c1").orderBy("value")

• partitionBy("c1")：指定按c1列分组，每组内独立计算排名
• orderBy("value")：指定组内按value列升序排序，默认排序规则为升序，如需降序可调整为orderBy(col("value").desc())，使用col需要额外导入from pyspark.sql.functions import col

3. 链式调用完成查询 #

对应你提供的SQL逻辑，完整DataFrame写法如下：

# 若已将tmp_view注册为Spark临时视图，可通过spark.table读取；如有对应DataFrame变量直接替换即可
result_df = spark.table("tmp_view") \
    # 新增rank列，计算密集排名
    .withColumn("rank", dense_rank().over(window_spec)) \
    # 过滤排名前3的数据
    .filter("rank <= 3")

如果需要查看结果，执行result_df.show()即可，输出结果和你提供的SQL执行结果完全一致。

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