python - Pyspark - 不确定如何将以下 X 行的总和分配给现有行值

Question

最好的解释方式是通过示例。在本例中，我们将采用接下来的两行:

原文:

ID  val
1   2
1   3    
1   1
1   9
2   1 
2   6
2   8
2   1

更新版本:

ID  sum_val
1   4
1   10    
1   9
1   0
2   14 
2   9
2   1
2   0

我正在 PySpark 中工作，因为我的数据集非常大。我是 PySpark 的新手，所以我在尝试完成这项工作时遇到了困难。

任何帮助将不胜感激。

Answer 1

使用窗口函数:

from pyspark.sql.functions import col, sum, monotonically_increasing_id
from pyspark.sql.window import Window

df = spark.createDataFrame(
    [(1, 2), (1, 3), (1, 1), (1, 9), (2, 1), (2, 6), (2, 8), (2, 1)],
    ("id", "val")
)

您需要像这样的Window:

w = (Window.partitionBy("id")
           .orderBy("_id")
           .rowsBetween(1, 2))

添加_id:

(df
   .withColumn("_id", monotonically_increasing_id())
   .withColumn("sum_val", sum("val").over(w))
   .na.fill(0)
   .show())

# +---+---+-----------+-------+          
# | id|val|        _id|sum_val|
# +---+---+-----------+-------+
# |  1|  2|          0|      4|
# |  1|  3|          1|     10|
# |  1|  1| 8589934592|      9|
# |  1|  9| 8589934593|      0|
# |  2|  1|17179869184|     14|
# |  2|  6|17179869185|      9|
# |  2|  8|25769803776|      1|
# |  2|  1|25769803777|      0|
# +---+---+-----------+-------+

请注意，这样的 monotonically_increasing_id 并不是一个好的做法 - 在生产中，您应该始终将排序信息嵌入到数据本身中，并且永远不要依赖 DataFrame< 的内部顺序.