spark使用总结--row_number使用和partial求多个字段均值

本文介绍Spark中的高级用法，包括使用窗口函数选取分组后特定条件的第一条数据，及利用functools.partial简化多个字段的均值计算并重命名，适用于大数据分析场景。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

最近项目中用到spark做简单的分析，总结下一些比较经典的用法。

一、窗口函数取分组后的第一条数据
第一条数据，可以是某个值最小最大等等，使用row_number和Window函数来实现。下面实列是获取每个班级里数学成绩最低的记录。其中，F.row_number().over(Window.partitionBy(“class”).orderBy(“math”))可以直接跟列名类似，在df的select中使用。

from pyspark.sql import functions as F

    select_cols = ["class"]
    cols = ["math"]
    select_cols.extend(cols)
    select_cols.append(F.row_number().over(Window.partitionBy("class").orderBy("math")).alias("rownum"))
    res_df = df.select(select_cols).filter("rownum == 1").select(select_cols[:-1])

二、partial实现对多个字段求均值且重命名

实现对dataframe求均值且重命名操作一般为：

res_df = df.groupby("g1").agg(
		F.avg("name1").alias("name1"),
		F.avg("name2").alias("name2"),
		F.avg("name3").alias("name3"),
		...)

这样操作，字段少的话直接写还可以接受，但是字段较多这样写就比较麻烦了，比如最近项目中的需要求均值的字段多达20个。
可以用functools.partial，输入列名数组和重命名数组，使用循环的方式来实现。

2.1 partial 简单说明
函数在执行时，要带上所有必要的参数进行调用。但是，有时参数可以在函数被调用之前提前获知。这种情况下，一个函数有一个或多个参数预先就能用上，以便函数能用更少的参数进行调用。
以下实例，重定义了函数add10，来设置默认的第一个参数为10。

import functools

def add(a, b, c):
    return a + b + c
if __name__ == "__main__":
    add10 = functools.partial(add, 10)
    print(add10(20, 30))

2.2 如何使用partial来实现
例如：对一个dataframe有字段：

字段	含义
class	班级
name	姓名
math	数学成绩
physics	物理成绩
chemistry	化学成绩

import functools
from pyspark.sql import functions as F

#对变量求均值  
select_cols = ["class"]
cols = ["math", "physics", "chemistry"]
select_cols.extend(cols)
select_df = df.select(select_cols)
gp_agg = functools.partial(select_df.groupby("class").agg)
for item in cols :
    gp_agg = functools.partial(gp_agg, F.avg(item).alias(item))
res_df = gp_agg()