Как перебирать каждую строку Dataframe/RDD в pyspark для группы.?

Question

Я хочу установить значение столбца на основе значения этого столбца в предыдущей строке для группы. Затем это обновленное значение будет использоваться в следующей строке.

У меня есть следующие dataframe

  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 

Группа имеет идентификатор и дату

Колонка C должна выводится на основе колонки А и В.

Если A == 1 и B == 0, то C получается C из предыдущей строки + 1.
Есть и другие условия, но я борюсь с этой частью.

Предполагая, что у нас есть столбец sort_date в dataframe.

Я попробовал следующий запрос:

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)) 

PS: Я использую Спарк 1.6

Answer 1

Сначала определим окно:

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

Ne xt создает UDAF, который получает A и B и в основном вычисляет C, начиная с 0, меняя значение на 0, когда появляется условие (A = 1, B = 0) и увеличивается на 1 в любое другое время. Чтобы увидеть, как написать UDAF увидеть примеры в here, here и here

EDIT Вот пример реализации в 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))