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

Question

I have a Spark DataFrame with ID and a bunch of numeric columns and for every column besides ID, I'm trying to generate a column of lagged differences grouped by ID.

For instance if I have this 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|
+---+-----+-----+-----+

I would expect the output to be something like this:

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

Where the _diff columns are the original columns subtracting their lags. My DataFrame has much more than 3 variables, so I'd want to be able to generate lagged differences for an arbitrarily high number of columns. ie I don't want to create the _diff columns one by one.

Any ideas as to how I can achieve this?

Answer 1

I would suggest you to go with foldLeft (a powerful api in 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)

which should give you

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

Note: I have used ID in orderBy which is not recomended, recomended to generate a separate column preserving the order of rows and use that instead of ID

I hope the answer is helpful

Answer 2

You will need to use lag as you pointed out in the function with Spark Window function.

You can generate a dynamic expression by storing the columns that you need to find the difference for.

The following will basically create an expression of type org.apache.spark.sql.Column which you can use over your original dataframe.

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

And then you can execute the expression generated above over your dataframe using 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|
+---+----+----+----+---------+---------+---------+