Question

我有像这种格式的患者数据

+---+-----+----+----------+
| id| name|code|      date|
+---+-----+----+----------+
|  1|Shaun|B121|2012-03-21|
|  3|Shaun|B120|2010-10-29|
|  2|Shaun|B121|2011-02-14|
|  4| John|B121|2011-09-29|
|  5| John|B120|2011-09-30|
|  6| John|B111|2012-09-30|
|  7| John|B121|2013-09-29|
+---+-----+----+----------+

我想检查代码为B121的每一行，如果历史代码B120是否应用于患者如果将set level设置为1，否则为2，对于代码为B120的行，将level 0设置为0。之后，结果应如下所示。

在MySQL中，我使用游标来做到这一点。

+---+-----+----+----------+-----+
| id| name|code|      date|level|
+---+-----+----+----------+-----+
|  3|Shaun|B120|2010-10-29|    0|
|  2|Shaun|B121|2011-02-14|    1|
|  1|Shaun|B121|2012-03-21|    1|
|  6| John|B111|2012-09-30|    0|
|  5| John|B120|2011-09-30|    0|
|  4| John|B121|2011-09-29|    2|
|  7| John|B121|2013-09-29|    1|
+---+-----+----+----------+-----+

已编辑：我添加了代码为B111的新行，但我想查看代码为B120的历史记录。

我试过这个解决方案

val window = Window.partitionBy("name").orderBy("date")
val lagCol = lag(col("date"), 1).over(window)

val pDF = df.withColumn("level", lagCol);

但它给出了以下结果

 id   name   code   date         level
 1    Shaun  B121   2012-03-21   2011-02-14
 2    Shaun  B121   2011-02-14   2010-10-19
 3    Shaun  B120   2010-10-19   Null
 5    John   B121   2013-09-29   2011-09-29
 4    John   B121   2011-09-29   Null

它检查前一行是代码B120还是B121，但我想用代码B120检查上一行。我不知道如何正确使用滞后功能。我是怎么做到的？

Answer 1

lag函数上的

Window函数会使您的要求变得复杂。

鉴于dataframe为

+---+-----+----+----------+
|id |name |code|date      |
+---+-----+----+----------+
|1  |Shaun|B121|2012-03-21|
|2  |Shaun|B121|2011-02-14|
|3  |Shaun|B120|2010-10-19|
|4  |John |B121|2011-09-29|
|5  |John |B121|2013-09-29|
+---+-----+----+----------+

您可以使用以下逻辑

import org.apache.spark.sql.functions._

val windowSpec = Window.partitionBy("name").orderBy("code", "date")

df.withColumn("temp", first("code").over(windowSpec))
    .withColumn("level", when($"temp" === "B120" && $"code" === "B121", 1).otherwise(when($"temp" === "B120" && $"code" === "B120", 0).otherwise(lit(2))))
    .drop("temp")

哪个应该给你

+---+-----+----+----------+-----+
|id |name |code|date      |level|
+---+-----+----+----------+-----+
|3  |Shaun|B120|2010-10-19|0    |
|2  |Shaun|B121|2011-02-14|1    |
|1  |Shaun|B121|2012-03-21|1    |
|4  |John |B121|2011-09-29|2    |
|5  |John |B121|2013-09-29|2    |
+---+-----+----+----------+-----+

已更新以满足其他要求

给出以下dataframe

+---+-----+----+----------+
|id |name |code|date      |
+---+-----+----+----------+
|1  |Shaun|B121|2012-03-21|
|2  |Shaun|B121|2011-02-14|
|3  |Shaun|B120|2010-10-29|
|4  |John |B121|2011-09-29|
|5  |John |B120|2011-09-30|
|6  |John |B111|2012-09-30|
|7  |John |B121|2013-09-29|
+---+-----+----+----------+

您可以创建一个udf函数来满足要求

import org.apache.spark.sql.functions._
def updateLevel = udf((array: mutable.WrappedArray[mutable.WrappedArray[String]]) => {
  val containsB120 = array.filter(ar => ar.contains("B120")).map(ar => (ar(1), ar(2)))
  var code = ""
  var date = "1970-01-01"
  if(containsB120.size > 0) {
    code = containsB120(0)._1
    date = containsB120(0)._2
  }
  val returnArray = array.map(row => {
    println(row(2), date, code)
    if(java.sql.Date.valueOf(row(2)).getTime > java.sql.Date.valueOf(date).getTime && code == "B120" && row(1) == "B121") {
      Array(row(0).toString, row(1).toString, row(2).toString, "1")
    }
    else if(java.sql.Date.valueOf(row(2)).getTime <= java.sql.Date.valueOf(date).getTime && row(1) == "B121" ) {
      Array(row(0).toString, row(1).toString, row(2).toString, "2")
    }
    else {
      Array(row(0).toString, row(1).toString, row(2).toString, "0")
    }
  })
  returnArray
})

然后在调用udf函数array和collect_list函数之前调用函数并在应用udf函数之后，应用explode函数来获取最终函数必需dataframe。

df.orderBy("date").withColumn("tempArray", array("id", "code", "date"))
    .groupBy("name")
    .agg(collect_list("tempArray").as("tempArray"))
    .withColumn("tempArray", explode(updateLevel($"tempArray")))
    .select($"tempArray"(0).as("id"), $"name", $"tempArray"(1).as("code"), $"tempArray"(2).as("date"), $"tempArray"(3).as("level"))

你应该有dataframe作为

+---+-----+----+----------+-----+
|id |name |code|date      |level|
+---+-----+----+----------+-----+
|3  |Shaun|B120|2010-10-29|0    |
|2  |Shaun|B121|2011-02-14|1    |
|1  |Shaun|B121|2012-03-21|1    |
|4  |John |B121|2011-09-29|2    |
|5  |John |B120|2011-09-30|0    |
|6  |John |B111|2012-09-30|0    |
|7  |John |B121|2013-09-29|1    |
+---+-----+----+----------+-----+

我希望答案很有帮助

Spark：在DataFrame中检查行的历史记录

1 个答案: