PySpark和SQL中的Window function

Window Function的使用

Window Function用于解决组内排名问题：

• 排名问题：每个部门按业绩来排名
• top N 问题：找出每个部门排名前 N 的员工进行奖励

from pyspark.sql.window import Window

# 创建window对象, 以相同uid为一个窗口
# 优先d降序, 其次ts降序
window = Window.partitionBy("uid") \
    .orderBy([F.col('d').desc(), F.col('ts').desc()])
    
# 在窗口内进行排序
df.withColumn("topn", F.row_number().over(window))

说明:

[100,100,98]的不同排序下的位次为
F.rank:113
F.dense_rank:112
F.row_number:123

代码例子

d1 = {'name1':["A","A","A","B","B","B",], 
      'age':[65,64,50,48,49,68],
      'height':[171,172,173,175,174,176,]
     }
df1 = spark.createDataFrame(pd.DataFrame(d1))
df1.createOrReplaceTempView("table1")
df1.show()
+-----+---+------+
|name1|age|height|
+-----+---+------+
|    A| 65|   171|
|    A| 64|   172|
|    A| 50|   173|
|    B| 48|   175|
|    B| 49|   174|
|    B| 68|   176|
+-----+---+------+

spark window function

from pyspark.sql.window import Window

window = Window.partitionBy("name1") \
    .orderBy([F.col('age').desc(), F.col('height').desc()])
    
# 在窗口内进行排序
df1.withColumn("topn", F.row_number().over(window))

df1.show()
+-----+---+------+----+
|name1|age|height|topn|
+-----+---+------+----+
|    A| 65|   171|   1|
|    A| 64|   172|   2|
|    A| 50|   173|   3|
|    B| 68|   176|   1|
|    B| 49|   174|   2|
|    B| 48|   175|   3|
+-----+---+------+----+

sql window function

spark.sql(
"""
select *, row_number()
over(partition by name1 
order by age desc) as topn
from table1
"""
).show()
+-----+---+------+----+
|name1|age|height|topn|
+-----+---+------+----+
|    A| 65|   171|   1|
|    A| 64|   172|   2|
|    A| 50|   173|   3|
|    B| 68|   176|   1|
|    B| 49|   174|   2|
|    B| 48|   175|   3|
+-----+---+------+----+