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的内部顺序