PySpark 数据框中类似 Pandas 的转换方式

在本文中，我们将介绍 PySpark 数据框中类似 Pandas 的一种转换方式。这种转换方式可以方便地对分组后的数据进行转换，类似于 Pandas 的 groupby+transform。

阅读更多：Pandas 教程

转换需要用到的三个函数

在 PySpark 中，我们需要使用三个函数来实现类似 Pandas 的 groupby+transform 转换，分别是 window、partitionBy 和 rowsBetween 函数。

window 函数用于定义窗口对应的列和偏移量，可以理解为分组的条件；
partitionBy 函数用于按照分组条件对数据进行分组；
rowsBetween 函数用于计算窗口的边界范围，用于计算汇总结果。

使用这三个函数，我们可以方便地对分组后的数据进行转换。

示例

下面我们以一个实际的示例来演示如何使用这三个函数实现类似 Pandas 的 groupby+transform 转换。

假设我们有一个包含以下字段的 PySpark 数据框：

+---+------+--------+
| id|group1|   value|
+---+------+--------+
|  1|     A|0.439263|
|  1|     B|0.756224|
|  1|     A|0.247512|
|  1|     B|0.318819|
|  2|     A|0.244810|
|  2|     B|0.793812|
|  2|     A|0.265070|
|  2|     B|0.700844|
|  3|     A|0.968261|
|  3|     B|0.767726|
+---+------+--------+

我们希望对其中的 value 列进行分组计算，按照 id 和 group1 字段进行分组，并计算每个分组中的最大值。则代码如下：

from pyspark.sql.functions import max, col, struct, lag
from pyspark.sql.window import Window

window_spec = Window.partitionBy("id", "group1").rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

df.withColumn("max_value", max(col("value")).over(window_spec)) \
  .show()

执行以上代码，可以得到以下结果：

+---+------+--------+----------+
| id|group1|   value| max_value|
+---+------+--------+----------+
|  1|     B|0.756224| 0.7562242|
|  1|     B|0.318819| 0.7562242|
|  1|     A|0.247512| 0.4392631|
|  1|     A|0.439263| 0.4392631|
|  2|     A|0.244810|  0.265070|
|  2|     A|0.265070|  0.265070|
|  2|     B|0.793812| 0.7938122|
|  2|     B|0.700844| 0.7938122|
|  3|     A|0.968261| 0.9682611|
|  3|     B|0.767726| 0.7677261|
+---+------+--------+----------+

上面的代码中，我们使用 window_spec 定义了一个窗口，只要是 id 和 group1 相同的行，都会被归为同一个窗口。通过 rowsBetween 函数，我们指定了对整个窗口进行汇总计算。

我们通过 withColumn 函数添加了一个新的列 max_value，其中使用了 max 函数对值进行计算，得到了在每个分组中的最大值。

总结

本文介绍了 PySpark 中一种类似 Pandas 的 groupby+transform 的方式，使用 window、partitionBy 和 rowsBetween 函数可以实现对分组后的数据进行汇总计算。通过实际示例演示，我们可以看到这种转换方式的便利性和实用性，我们可以在 PySpark 中更加方便地对数据进行处理和转换，进而实现各种数据分析和挖掘任务。