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創建_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

正如您在帶有 Spark Window 函數的函數中指出的那樣，您將需要使用lag 。

您可以通過存儲需要查找差異的列來生成動態表達式。

下面將基本上創建一個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|
+---+----+----+----+---------+---------+---------+