Apache Spark - 如何为数据框中的每一列创建不同的列？

Question

我有一个带有 ID 和一堆数字列的 Spark DataFrame，对于除 ID 之外的每一列，我正在尝试生成一列按 ID 分组的滞后差异。

例如，如果我有这个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

我建议你使用foldLeft（scala 中一个强大的 api）

//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        |
+---+----+----+----+---------+---------+---------+

注意：我在 orderBy 中使用了不推荐的 ID，建议生成一个单独的列来保留行的顺序并使用它来代替 ID

希望回答对你有帮助

Answer 2

您需要使用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|
+---+----+----+----+---------+---------+---------+