Scala Dataframe窗口滞后功能条件答案

【问题标题】：Scala Dataframe window lag function on conditionScala Dataframe窗口滞后功能条件
【发布时间】：2018-06-19 16:18:15
【问题描述】：

我有一个包含以下记录的数据框。

id   col1         time
A      1    2017-12-23 09:00
A-1    1    2017-12-23 11:00
A-1    1    2017-12-23 10:00
A      2    2017-12-23 08:00
A      3    2017-12-23 07:00
A-2    4    2017-12-23 12:00
A-1    4    2017-12-23 12:00
B      1    2017-12-23 09:00
B-1    1    2017-12-23 11:00
B-1    1    2017-12-23 10:00
B      2    2017-12-23 08:00
B      3    2017-12-23 07:00
B-2    4    2017-12-23 12:00
B-1    4    2017-12-23 12:00

我想从上述数据框中获取滞后特征，但在“id”列上有一个条件。是否可以在不更改记录的情况下给出滞后特征的条件？像下面这样累了

val window = Window.partitionBy("id").orderBy("time")
val resultDF = DF.withColumn("col1_lag",lag(DF("col1"), 1).over(window))

上面的代码将导致基于分组 id 的滞后特征，我想考虑带有“-”的值是孩子的，他们需要父母的滞后特征。例如：id 为 A-1(child) 的行需要 A(parent) 的值出现在它之前。

预期的结果应该如下所示。

id   col1         time           col1_lag
A      3    2017-12-23 07:00       null
A      2    2017-12-23 08:00       3
A      1    2017-12-23 09:00       2
A-1    1    2017-12-23 10:00       1   
A-1    1    2017-12-23 11:00       1  
A-1    4    2017-12-23 12:00       1
A-2    4    2017-12-23 12:00       1
B      3    2017-12-23 07:00       null 
B      2    2017-12-23 08:00       3   
B      1    2017-12-23 09:00       2
B-1    1    2017-12-23 10:00       1
B-1    1    2017-12-23 11:00       1 
B-1    4    2017-12-23 12:00       1
B-2    4    2017-12-23 12:00       1

【问题讨论】：

标签： scala dataframe spark-dataframe lag feature-extraction

【解决方案1】：

如果我理解您的要求更正，则需要相当多的步骤才能实现您的要求（因此成本相当高）：

扩展 DF 以包括识别 row id（如果 DF 已经有则不需要）、parent id、isChild 和 lag(col1) over Window partition byparent id`
通过parent id 和isChild 扩展DF，保留row id 和col1 lag 的列表
调整child 行的col1 lag 列表（即isChild == true）以使所有元素与列表的第一个元素相同，通过UDF 使用row id 压缩它，然后重新展开分组的DF
通过row id 将第 1 步中的扩展 DF 与第 3 步中调整的 DF 连接起来，以获得调整后的col1 lag

示例代码如下：

val DF = Seq(
  ("A", 1, "2017-12-23 09:00"),
  ("A-1", 1, "2017-12-23 11:00"),
  ("A-1", 1, "2017-12-23 10:00"),
  ("A", 2, "2017-12-23 08:00"),
  ("A", 3, "2017-12-23 07:00"),
  ("A-2", 4, "2017-12-23 12:00"),
  ("A-1", 4, "2017-12-23 12:00"),
  ("B", 1, "2017-12-23 09:00"),
  ("B-1", 1, "2017-12-23 11:00"),
  ("B-1", 1, "2017-12-23 10:00"),
  ("B", 2, "2017-12-23 08:00"),
  ("B", 3, "2017-12-23 07:00"),
  ("B-2", 4, "2017-12-23 12:00"),
  ("B-1", 4, "2017-12-23 12:00")
).toDF("id", "col1", "time")

import org.apache.spark.sql.expressions.Window
val winPid = Window.partitionBy("pid").orderBy("id", "time")

// Step 1 
val expandedDF = DF.
  withColumn("row_id", monotonically_increasing_id).
  withColumn("pid", substring($"id", 0, 1)).
  withColumn("is_child", when($"pid" === $"id", false).otherwise(true)).
  withColumn("col1_lag", lag($"col1", 1, 0).over(winPid)).
  orderBy($"pid", $"id", $"time")

expandedDF.show
// +---+----+----------------+------+---+--------+--------+
// | id|col1|            time|row_id|pid|is_child|col1_lag|
// +---+----+----------------+------+---+--------+--------+
// |  A|   3|2017-12-23 07:00|     4|  A|   false|       0|
// |  A|   2|2017-12-23 08:00|     3|  A|   false|       3|
// |  A|   1|2017-12-23 09:00|     0|  A|   false|       2|
// |A-1|   1|2017-12-23 10:00|     2|  A|    true|       1|
// |A-1|   1|2017-12-23 11:00|     1|  A|    true|       1|
// |A-1|   4|2017-12-23 12:00|     6|  A|    true|       1|
// |A-2|   4|2017-12-23 12:00|     5|  A|    true|       4|
// |  B|   3|2017-12-23 07:00|    11|  B|   false|       0|
// |  B|   2|2017-12-23 08:00|    10|  B|   false|       3|
// |  B|   1|2017-12-23 09:00|     7|  B|   false|       2|
// |B-1|   1|2017-12-23 10:00|     9|  B|    true|       1|
// |B-1|   1|2017-12-23 11:00|     8|  B|    true|       1|
// |B-1|   4|2017-12-23 12:00|    13|  B|    true|       1|
// |B-2|   4|2017-12-23 12:00|    12|  B|    true|       4|
// +---+----+----------------+------+---+--------+--------+

// Step 2
val groupedDF = expandedDF.groupBy("pid", "is_child").agg(
  collect_list("row_id").as("row_id_list"),
  collect_list("col1_lag").as("col1_lag_list")
).orderBy($"pid", $"is_child")

groupedDF.show
// +---+--------+--------------+-------------+
// |pid|is_child|   row_id_list|col1_lag_list|
// +---+--------+--------------+-------------+
// |  A|   false|     [4, 3, 0]|    [0, 3, 2]|
// |  A|    true|  [2, 1, 6, 5]| [1, 1, 1, 4]|
// |  B|   false|   [11, 10, 7]|    [0, 3, 2]|
// |  B|    true|[9, 8, 13, 12]| [1, 1, 1, 4]|
// +---+--------+--------------+-------------+

// Step 3
def adjustAndZip = udf(
  (row: Seq[Long], lag: Seq[Int], isChild: Boolean) => {
    val adjustedLag = if (isChild) Seq.fill[Int](lag.length)(lag.head) else lag
    row zip adjustedLag
  })

val adjustedDF = groupedDF.withColumn("temp",
  explode(adjustAndZip($"row_id_list", $"col1_lag_list", $"is_child"))
).select(
  $"temp._1".as("row_id"), $"temp._2".as("col1_lag_adjusted")
)

adjustedDF.show
// +------+-----------------+
// |row_id|col1_lag_adjusted|
// +------+-----------------+
// |     4|                0|
// |     3|                3|
// |     0|                2|
// |     2|                1|
// |     1|                1|
// |     6|                1|
// |     5|                1|
// |    11|                0|
// |    10|                3|
// |     7|                2|
// |     9|                1|
// |     8|                1|
// |    13|                1|
// |    12|                1|
// +------+-----------------+

// Step 4
val resultDF = expandedDF.join(adjustedDF, Seq("row_id")).select(
  $"pid", $"id", $"col1", $"time", $"col1_lag_adjusted"
).orderBy($"pid", $"id", $"time")

resultDF.show
// +---+---+----+----------------+-----------------+
// |pid| id|col1|            time|col1_lag_adjusted|
// +---+---+----+----------------+-----------------+
// |  A|  A|   3|2017-12-23 07:00|                0|
// |  A|  A|   2|2017-12-23 08:00|                3|
// |  A|  A|   1|2017-12-23 09:00|                2|
// |  A|A-1|   1|2017-12-23 10:00|                1|
// |  A|A-1|   1|2017-12-23 11:00|                1|
// |  A|A-1|   4|2017-12-23 12:00|                1|
// |  A|A-2|   4|2017-12-23 12:00|                1|
// |  B|  B|   3|2017-12-23 07:00|                0|
// |  B|  B|   2|2017-12-23 08:00|                3|
// |  B|  B|   1|2017-12-23 09:00|                2|
// |  B|B-1|   1|2017-12-23 10:00|                1|
// |  B|B-1|   1|2017-12-23 11:00|                1|
// |  B|B-1|   4|2017-12-23 12:00|                1|
// |  B|B-2|   4|2017-12-23 12:00|                1|
// +---+---+----+----------------+-----------------+

【讨论】：

如果有价值，什么是昂贵的？