Question

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

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（子）的行需要A（父）的值出现在它之前。

预期结果应如下所示。

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

Answer 1

如果我了解您的需求更正，那么需要采取相当多的步骤来实现您的要求（因此成本相当高）：

展开DF以包含标识row id（如果DF已有，则不需要），parent id，isChild和lag(col1) over Window partition by父ID
按parent id和isChild扩展了DF，保留了row id和col1 lag
调整col1 lag行的child列表（即isChild == true），使所有元素与列表的第一个元素相同，用row id压缩通过UDF，重新扩展分组DF
在步骤3中加入展开的DF，并在步骤3中按row id加入已调整的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|
// +---+---+----+----------------+-----------------+

条件下Scala Dataframe窗口滞后功能

1 个答案: