python - 使用 Spark 对分组数据进行快速行间计算

Question

我有一个具有以下结构的 CSV 文件(> 3GB)

groupId | time | value
1            0       2 
1            1       1 
1            2       4 
2            0       6 
2            1       2 

并且想要添加一列(值 t-1)，其中包含同一组中早一个时间步的“行”的值:

groupId | time | value | value t-1
1            0       2           -
1            1       1           2
1            2       4           1
2            0       6           -
2            1       2           6

我想成本高昂的部分是搜索前一行。不知何故，这似乎是一项没有缩减的 map 缩减工作——如果这有意义的话。但据我了解，我不能确定同一组的所有数据都在同一个工作人员上。

Spark 是适合这项工作的工具吗？

我最好的替代解决方案是将文件拆分为多个文件(每组一个)，然后运行 ​​python 脚本的多个实例，对值 t-1 进行排序并按顺序添加值

Answer 1

这可以通过使用 Spark 的 Window 函数来实现，如下所示。

import org.apache.spark.sql.expressions.Window

val df = Seq((1,0,2), (1,1,1), (1,2,4), (2,0,6), (2,1,2)).toDF("groupId", "time", "value")

val result = df.withColumn("value_t-1", sum($"value").over(Window.partitionBy("groupId").orderBy("time").rowsBetween(-1,-1)))

输出:

scala> result.show()
+-------+----+-----+---------+
|groupId|time|value|value_t-1|
+-------+----+-----+---------+
|      1|   0|    2|     null|
|      1|   1|    1|        2|
|      1|   2|    4|        1|
|      2|   0|    6|     null|
|      2|   1|    2|        6|
+-------+----+-----+---------+

Python 版本:

>>> from pyspark.sql.window import Window
>>> import pyspark.sql.functions as func
>>> df = spark.createDataFrame([(1,0,2), (1,1,1), (1,2,4), (2,0,6), (2,1,2)], ["groupId", "time", "value"])
>>> result = df.withColumn("value_t-1", func.sum(df.value).over(Window.partitionBy(df.groupId).orderBy(df.time).rowsBetween(-1,-1)))
>>> result.show()
+-------+----+-----+---------+
|groupId|time|value|value_t-1|
+-------+----+-----+---------+
|      1|   0|    2|     null|
|      1|   1|    1|        2|
|      1|   2|    4|        1|
|      2|   0|    6|     null|
|      2|   1|    2|        6|
+-------+----+-----+---------+