Apache Spark - как создать разностные столбцы для каждого столбца в датафрейме?

Question

У меня есть Spark DataFrame с идентификатором и набором числовых столбцов, и для каждого столбца, кроме идентификатора, я пытаюсь сгенерировать столбец лаговых различий, сгруппированных по идентификатору.

Например, если у меня есть это DataFrame

+---+-----+-----+-----+
| ID| var1| var2| var3|
+---+-----+-----+-----+
|  1|    1|    3|    2|
|  1|    2|    4|    2|
|  1|    3|    1|    3|
|  2|    1|    3|    4|
|  2|    1|    2|    1|
|  2|    1|    1|    1|
|  2|    3|    3|    1|
|  3|   -1|    0|    0|
|  3|    2|   -1|    2|
|  3|    0|    4|    0|
+---+-----+-----+-----+

Я ожидаю, что результат будет примерно таким:

+---+-----+-----+-----+----------+----------+----------+
| ID| var1| var2| var3| var1_diff| var2_diff| var3_diff|
+---+-----+-----+-----+----------+----------+----------+
|  1|    1|    3|    2|      null|      null|      null|
|  1|    2|    4|    2|         1|         1|         0|
|  1|    3|    1|    3|         1|        -3|         1|
|  2|    1|    3|    4|      null|      null|      null|
|  2|    1|    2|    1|         0|        -1|        -3|
|  2|    1|    1|    1|         0|        -1|         0|
|  2|    3|    3|    1|         2|         2|         0|
|  3|   -1|    0|    0|      null|      null|      null|
|  3|    2|   -1|    2|         3|        -1|         2|
|  3|    0|    4|    0|        -2|         3|        -2|
+---+-----+-----+-----+----------+----------+----------+

Где столбцы _diffисходные столбцы, вычитающие их лагиМой DataFrame имеет более 3 переменных, поэтому я хотел бы иметь возможность генерировать лаговые различия для произвольно большого числа столбцов.то есть я не хочу создавать столбцы _diff по одному.

Есть идеи, как мне этого добиться?

Answer 1

Вам нужно будет использовать lag, как вы указали в функции с функцией Spark Window.

Вы можете сгенерировать динамическое выражение, сохранив столбцы, для которых нужно найти разницу.

Следующее в основном создаст выражение типа org.apache.spark.sql.Column, которое вы можете использовать в исходном фрейме данных.

import org.apache.spark.sql.expressions.Window
val w = Window.partitionBy($"id").orderBy($"id")

//df.columns returns all the columns of the dataframe
//union is used to include the original columns in the expression
//expr looks like : (var1 - lag(var1) over window) as var1_diff ...

val expr = df.columns.map(col(_)) union df.columns.filterNot(_.toLowerCase.equals("id")).map { x => (col(x) - lag(col(x),1).over(w) ).as(s"${x}_diff") }

И затем вы можете выполнить сгенерированное выше выражение для вашего фрейма данных, используя select

df.select(expr:_*).show
+---+----+----+----+---------+---------+---------+
| id|var1|var2|var3|var1_diff|var2_diff|var3_diff|
+---+----+----+----+---------+---------+---------+
|  1|   1|   3|   2|     null|     null|     null|
|  1|   2|   4|   2|        1|        1|        0|
|  1|   3|   1|   3|        1|       -3|        1|
|  3|  -1|   0|   0|     null|     null|     null|
|  3|   2|  -1|   2|        3|       -1|        2|
|  3|   0|   4|   0|       -2|        5|       -2|
|  2|   1|   3|   4|     null|     null|     null|
|  2|   1|   2|   1|        0|       -1|       -3|
|  2|   1|   1|   1|        0|       -1|        0|
|  2|   3|   3|   1|        2|        2|        0|
+---+----+----+----+---------+---------+---------+

Answer 2

Я бы предложил вам использовать foldLeft (мощный API в Scala)

//assuming that the column ID is at the front
val tailColumns = df.columns.tail

import org.apache.spark.sql.expressions._
def windowSpec = Window.partitionBy("ID").orderBy("ID")

tailColumns.foldLeft(df){(tempdf, colName) => tempdf.withColumn(colName+"_diff", col(colName)-lag(col(colName), 1).over(windowSpec))}.show(false)

, который должен дать вам

+---+----+----+----+---------+---------+---------+
|ID |var1|var2|var3|var1_diff|var2_diff|var3_diff|
+---+----+----+----+---------+---------+---------+
|1  |1   |3   |2   |null     |null     |null     |
|1  |2   |4   |2   |1        |1        |0        |
|1  |3   |1   |3   |1        |-3       |1        |
|3  |-1  |0   |0   |null     |null     |null     |
|3  |2   |-1  |2   |3        |-1       |2        |
|3  |0   |4   |0   |-2       |5        |-2       |
|2  |1   |3   |4   |null     |null     |null     |
|2  |1   |2   |1   |0        |-1       |-3       |
|2  |1   |1   |1   |0        |-1       |0        |
|2  |3   |3   |1   |2        |2        |0        |
+---+----+----+----+---------+---------+---------+

Примечание: я использовалИдентификатор в порядке, который не рекомендуется, рекомендуется создавать отдельный столбец с сохранением порядка строк и использовать его вместо идентификатора

Надеюсь, ответ полезен