PySpark Window.partitionBy 和 groupBy 两个函数的区别和使用场景

在本文中，我们将介绍 PySpark 中的 Window.partitionBy 和 groupBy 两个函数的区别和使用场景。

阅读更多：PySpark 教程

Window.partitionBy

Window.partitionBy 函数是 PySpark 中的窗口函数，用于对数据进行分组操作，类似于 SQL 中的 PARTITION BY 子句。它可以将数据集按照指定的列进行分组，并为每个分组创建一个子集。在窗口函数之前使用 Window.partitionBy，可以根据指定的列对数据进行分组，然后在每个分组内进行计算。

示例：

假设我们有一个包含以下列的数据集：name、age、gender和salary。我们想要按照 gender 列对数据进行分组，并计算每个分组中的平均薪资。

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import avg

spark = SparkSession.builder.getOrCreate()

# 读取数据集
df = spark.read.csv("data.csv", header=True, inferSchema=True)

# 定义窗口规范
windowSpec = Window.partitionBy("gender")

# 使用窗口函数计算每个分组的平均薪资
df.withColumn("avg_salary", avg("salary").over(windowSpec)).show()

上述代码中，通过 Window.partitionBy(“gender”) 指定了按照 gender 列进行分组，然后使用 avg(“salary”).over(windowSpec) 计算了每个分组内 salary 列的平均值，并将结果添加到了新的一列 “avg_salary” 中。

groupBy

groupBy 函数是 PySpark 中的聚合函数，用于对数据进行分组操作，类似于 SQL 中的 GROUP BY 子句。它可以将数据集按照指定的列进行分组，并对每个分组进行聚合计算。与 Window.partitionBy 不同，groupBy 返回的是一个新的 DataFrame，而不是在原始 DataFrame 中添加新的列。

示例：

假设我们有一个包含以下列的数据集：name、age、gender和salary。我们想要按照 gender 列对数据进行分组，并计算每个分组中的最高薪资。

from pyspark.sql import SparkSession
from pyspark.sql.functions import max

spark = SparkSession.builder.getOrCreate()

# 读取数据集
df = spark.read.csv("data.csv", header=True, inferSchema=True)

# 使用 groupBy 函数计算每个分组的最高薪资
df.groupBy("gender").agg(max("salary").alias("max_salary")).show()

上述代码中，通过 groupBy(“gender”) 指定了按照 gender 列进行分组，然后使用 agg(max(“salary”).alias(“max_salary”)) 计算了每个分组内 salary 列的最大值，并将结果添加到了新的一列 “max_salary” 中。

使用场景

Window.partitionBy 适用于需要在每个分组内进行复杂的计算或对结果进行排序的场景。通过窗口函数，我们可以对每个分组内的数据进行聚合、排序、计数等操作，并将结果作为新的列添加到原始 DataFrame 中。

groupBy 适用于只需要对每个分组进行基本聚合计算的场景。groupBy 返回的是一个新的 DataFrame，可以使用聚合函数对每个分组进行计算，并进行汇总统计。

总结

本文介绍了 PySpark 中的 Window.partitionBy 和 groupBy 两个函数的区别和使用场景。Window.partitionBy 适用于需要在每个分组内进行复杂的计算或排序的场景，可以通过窗口函数对每个分组内的数据进行聚合操作。groupBy 适用于只需要对每个分组进行基本聚合计算的场景，返回的是一个新的 DataFrame。根据具体的需求，我们可以选择合适的函数进行数据分组和聚合操作。