如何在 PySpark 中为一个组迭代 Dataframe / RDD 的每一行。？

Posted 2023-03-31

【中文标题】如何在 PySpark 中为一个组迭代 Dataframe / RDD 的每一行。？

【英文标题】：How to iterate over each row of an Dataframe / RDD in PySpark for a group.?

【发布时间】：2017-01-30 06:07:53

【问题描述】：

我想根据组的前一行中该列的值设置列的值。然后这个更新后的值将用于下一行。

我有以下数据框

id | start_date|sort_date | A | B |
-----------------------------------
1 | 1/1/2017 | 31-01-2015 | 1 | 0 | 
1 | 1/1/2017 | 28-02-2015 | 0 | 0 | 
1 | 1/1/2017 | 31-03-2015 | 1 | 0 | 
1 | 1/1/2017 | 30-04-2015 | 1 | 0 | 
1 | 1/1/2017 | 31-05-2015 | 1 | 0 | 
1 | 1/1/2017 | 30-06-2015 | 1 | 0 | 
1 | 1/1/2017 | 31-07-2015 | 1 | 0 | 
1 | 1/1/2017 | 31-08-2015 | 1 | 0 | 
1 | 1/1/2017 | 30-09-2015 | 0 | 0 | 
2 | 1/1/2017 | 31-10-2015 | 1 | 0 | 
2 | 1/1/2017 | 30-11-2015 | 0 | 0 | 
2 | 1/1/2017 | 31-12-2015 | 1 | 0 | 
2 | 1/1/2017 | 31-01-2016 | 1 | 0 | 
2 | 1/1/2017 | 28-02-2016 | 1 | 0 | 
2 | 1/1/2017 | 31-03-2016 | 1 | 0 | 
2 | 1/1/2017 | 30-04-2016 | 1 | 0 | 
2 | 1/1/2017 | 31-05-2016 | 1 | 0 | 
2 | 1/1/2017 | 30-06-2016 | 0 | 0 | 

输出：

id | start_date|sort_date | A | B | C
---------------------------------------
1 | 1/1/2017 | 31-01-2015 | 1 | 0 | 1
1 | 1/1/2017 | 28-02-2015 | 0 | 0 | 0
1 | 1/1/2017 | 31-03-2015 | 1 | 0 | 1
1 | 1/1/2017 | 30-04-2015 | 1 | 0 | 2
1 | 1/1/2017 | 31-05-2015 | 1 | 0 | 3
1 | 1/1/2017 | 30-06-2015 | 1 | 0 | 4
1 | 1/1/2017 | 31-07-2015 | 1 | 0 | 5
1 | 1/1/2017 | 31-08-2015 | 1 | 0 | 6
1 | 1/1/2017 | 30-09-2015 | 0 | 0 | 0
2 | 1/1/2017 | 31-10-2015 | 1 | 0 | 1
2 | 1/1/2017 | 30-11-2015 | 0 | 0 | 0
2 | 1/1/2017 | 31-12-2015 | 1 | 0 | 1
2 | 1/1/2017 | 31-01-2016 | 1 | 0 | 2
2 | 1/1/2017 | 28-02-2016 | 1 | 0 | 3
2 | 1/1/2017 | 31-03-2016 | 1 | 0 | 4
2 | 1/1/2017 | 30-04-2016 | 1 | 0 | 5
2 | 1/1/2017 | 31-05-2016 | 1 | 0 | 6
2 | 1/1/2017 | 30-06-2016 | 0 | 0 | 0

组的 id 和日期

C 列是根据 A 列和 B 列推导出来的。

如果 A == 1 且 B == 0，则 C 是从前一行 + 1 派生的 C。 还有一些其他条件，但我正在努力解决这部分问题。

假设我们在数据框中有一列 sort_date。

我尝试了以下查询：

SELECT
id,
date,
sort_date,
lag(A) OVER (PARTITION BY  id, date ORDER BY sort_date) as prev,
CASE
   WHEN A=1 AND B= 0  THEN 1
   WHEN  A=1 AND B> 0 THEN prev +1
   ELSE 0
 END AS A
FROM
Table

这就是我为 UDAF 所做的事情

val myFunc = new MyUDAF
val w = Window.partitionBy(col("ID"), col("START_DATE")).orderBy(col("SORT_DATE"))
val df = df.withColumn("C", myFunc(col("START_DATE"), col("X"),
  col("Y"), col("A"),
  col("B")).over(w))

P.S : 我使用的是 Spark 1.6

【问题讨论】：

您可以在 Spark SQL 中使用 窗口函数。

你可以添加你尝试过的代码吗？

请改进问题：你能多解释一下你想要达到的目标吗，到目前为止你做了什么，你的输入是什么，你的预期输出是什么，你想这样做吗？标题所说的RDD还是列的措辞所暗示的数据框？你是什​​么意思组？你的意思是groupby？你希望它如何排序？

抱歉，我是通过手机执行此操作的。我已经编辑了这个问题。我正在使用数据框。还有一个日期列，用于对组中的行进行排序。

添加你迄今为止尝试过的东西也很好。并且还要回答其他 cmets，因为您的问题非常广泛并且也可能被关闭。

【参考方案1】：

首先定义一个窗口：

import org.apache.spark.sql.expressions.Window
val winspec = Window.partitionBy("id","start_date").orderBy("sort_date")

接下来创建一个接收 A 和 B 的 UDAF，基本上从 0 开始计算 C，每当条件出现时更改为 0 (A=1,B=0) 并在任何其他时间增加 1。要了解如何编写 UDAF，请参阅 here、here 和 here 中的示例

编辑 这是 UDAF 的示例实现（未经过真正测试，因此可能存在拼写错误）：

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.MutableAggregationBuffer,UserDefinedAggregateFunction
 import org.apache.spark.sql.types._

 class myFunc() extends UserDefinedAggregateFunction 

  // Input Data Type Schema
  def inputSchema: StructType = StructType(Array(StructField("A", IntegerType), StructField("A", IntegerType)))

   // Intermediate Schema
  def bufferSchema = StructType(Array(StructField("C", IntegerType)))

  // Returned Data Type .
  def dataType: DataType = IntegerType

  // Self-explaining
  def deterministic = true

  // This function is called whenever key changes
  def initialize(buffer: MutableAggregationBuffer) = 
    buffer(0) = 0 // set number of items to 0
  

  // Iterate over each entry of a group
  def update(buffer: MutableAggregationBuffer, input: Row) = 
    buffer(0) = if (input.getInt(0) == 1 && input.getInt(1) == 0) buffer.getInt(0) + 1 else 0
  

  // Merge two partial aggregates - doesn't really matter because the window will make sure the buffer remains in a
  // single partition
  def merge(buffer1: MutableAggregationBuffer, buffer2: Row) = 
    buffer1(0) = buffer1.getInt(0) + buffer2.getInt(0)
  

  // Called after all the entries are exhausted.
  def evaluate(buffer: Row) = 
    buffer.getInt(0)
  

最后将其应用于您的数据框。假设您将 UDAF 命名为 myFunc：

val f = new myFunc()
val newDF = df.withColumn("newC", f($"A",$"B").over(winspec))

【讨论】：

C 是派生列。最初每一行都有 C=0。我需要根据前一行计算的 C 值来计算当前行的 C 。在这种情况下，前一个 C 将始终为 0。

我已更改我的问题以反映这一点。这是我的错。

这里我只是叫它newC。创建计算的部分是 UDAF。 UDAF 将以 0 开始 newC 并将其增加 1，除非 A=1 和 B=0。窗口所做的是确保对 UDAF 的输入进行正确分区和排序

如何增加 newC. newC 的先前值是否可用于下一行？

def myFunc(A,B): if A == 1 and B == 0: return 1 else: 将前一个值增加 1 () # 只是一个例子