基于另一列的新列和 spark 中的值更改

Question

我有一个火花 dataframe ，我想根据列ID和DAT_TIMESTAMP创建一个新列RESULT ，具体取决于STATUS列的值。 当STATUS值第一次变为A时，我想获取DAT_TIMESTAMP的值并将其与ID的值连接，并继续对后续行使用相同的DAT_TIMESTAMP值，直到STATUS再次返回A

我希望下面的示例数据有助于解释我想要做什么。

+--------+--------------+-------+----------------------+
|      ID| DAT_TIMESTAMP| STATUS|                RESULT|
+--------+--------------+-------+----------------------+
| ID_1111| 1617214599502|      D| ID_1111_1617214600304|
| ID_1111| 1617214600002|      D| ID_1111_1617214600304|
| ID_1111| 1617214600502|      A| ID_1111_1617214600502| // first appearance of A
| ID_1111| 1617214601003|      A| ID_1111_1617214600502|
| ID_1111| 1617214601503|      A| ID_1111_1617214600502|
| ID_1111| 1617214602003|      B| ID_1111_1617214600502|
| ID_1111| 1617214602503|      B| ID_1111_1617214600502|
| ID_1111| 1617214603004|      C| ID_1111_1617214600502|
| ID_1111| 1617214603504|      C| ID_1111_1617214600502|
| ID_1111| 1617214604004|      C| ID_1111_1617214600502|
| ID_1111| 1617214604504|      C| ID_1111_1617214600502|
| ID_1111| 1617214605003|      D| ID_1111_1617214600502|
| ID_1111| 1617214605506|      D| ID_1111_1617214600502|
| ID_1111| 1617214606003|      A| ID_1111_1617214606003| // first appearance of A, again
| ID_1111| 1617214606504|      A| ID_1111_1617214606003|
| ID_1111| 1617214607004|      A| ID_1111_1617214606003|

我一直在尝试when function 来做到这一点，但我真的不确定我在做什么。 这可以在火花中完成吗？ 谢谢。

Answer 1

终于做到了，也许不是最优雅或最好的方法......

我可以通过使用lag function 创建一个包含先前STATUS值的中间列来确认STATUS列中的值是否第一次变为A 然后比较当前值是否为A和先前的值不是A 。

import org.apache.spark.sql.expressions.Window

// load data...


// add new column
val df2 = df.orderBy("ID", "DAT_TIMESTAMP")
            .withColumn("RESULT",
                last(
                    when(
                        $"STATUS"==="A" &&

                        // intermediate column for comparison
                        lag($"STATUS", 1).over(Window.orderBy("ID", "DAT_TIMESTAMP")).notEqual("A"),

                        concat_ws("_", $"ID", $"DAT_TIMESTAMP"))
                    .otherwise(null), true)
                .over(Window.orderBy("ID", "DAT_TIMESTAMP")))
            .select("*")

df2 如下所示。 RESULT的前两行是null因为以前没有对A引用。

df2.show(false)

+-------+-------------+------+---------------------+
|ID     |DAT_TIMESTAMP|STATUS|RESULT               |
+-------+-------------+------+---------------------+
|ID_1111|1617214599502|D     |null                 |
|ID_1111|1617214600002|D     |null                 |
|ID_1111|1617214600502|A     |ID_1111_1617214600502|
|ID_1111|1617214601003|A     |ID_1111_1617214600502|
|ID_1111|1617214601503|A     |ID_1111_1617214600502|
|ID_1111|1617214602003|B     |ID_1111_1617214600502|
|ID_1111|1617214602503|B     |ID_1111_1617214600502|
|ID_1111|1617214603004|C     |ID_1111_1617214600502|
|ID_1111|1617214603504|C     |ID_1111_1617214600502|
|ID_1111|1617214604004|C     |ID_1111_1617214600502|
|ID_1111|1617214604504|C     |ID_1111_1617214600502|
|ID_1111|1617214605003|D     |ID_1111_1617214600502|
|ID_1111|1617214605506|D     |ID_1111_1617214600502|
|ID_1111|1617214606003|A     |ID_1111_1617214606003|
|ID_1111|1617214606504|A     |ID_1111_1617214606003|
|ID_1111|1617214607004|A     |ID_1111_1617214606003|
+-------+-------------+------+---------------------+