数据框获取对应列的第一个和最后一个值

Question

Is it possible to get first value of the corresponding column within subgroup.

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.expressions.{Window, WindowSpec}

object tmp {
  def main(args: Array[String]): Unit = {
    val spark =  SparkSession.builder().master("local").getOrCreate()
    import spark.implicits._

    val input = Seq(
      (1235,  1, 1101, 0),
      (1235,  2, 1102, 0),
      (1235,  3, 1103, 1),
      (1235,  4, 1104, 1),
      (1235,  5, 1105, 0),
      (1235,  6, 1106, 0),
      (1235,  7, 1107, 1),
      (1235,  8, 1108, 1),
      (1235,  9, 1109, 1),
      (1235, 10, 1110, 0),
      (1235, 11, 1111, 0)
    ).toDF("SERVICE_ID", "COUNTER", "EVENT_ID", "FLAG")

    lazy val window: WindowSpec = Window.partitionBy("SERVICE_ID").orderBy("COUNTER")
    val firsts = input.withColumn("first_value", first("EVENT_ID", ignoreNulls = true).over(window.rangeBetween(Long.MinValue, Long.MaxValue)))
    firsts.orderBy("SERVICE_ID", "COUNTER").show()

  }
}

Output I want.

First (or Previous) value of column EVENT_ID based on FLAG = 1 And Last (or Next ) value of column EVENT_ID based on FLAG = 1 partition by SERVICE_ID sorted by counter

+----------+-------+--------+----+-----------+-----------+
|SERVICE_ID|COUNTER|EVENT_ID|FLAG|first_value|last_value|
+----------+-------+--------+----+-----------+-----------+
|      1235|      1|    1101|   0|          0|       1103|
|      1235|      2|    1102|   0|          0|       1103|
|      1235|      3|    1103|   1|          0|       1106|
|      1235|      4|    1104|   0|       1103|       1106|
|      1235|      5|    1105|   0|       1103|       1106|
|      1235|      6|    1106|   1|          0|       1108|
|      1235|      7|    1107|   0|       1106|       1108|
|      1235|      8|    1108|   1|          0|       1109|
|      1235|      9|    1109|   1|          0|       1110|
|      1235|     10|    1110|   1|          0|          0|
|      1235|     11|    1111|   0|       1110|          0|
|      1235|     12|    1112|   0|       1110|          0|
+----------+-------+--------+----+-----------+-----------+

Answer 1

First the dataframe need to be formed into groups. A new group starts at each time the "TIME" column equals 1. To do this, first add a column "ID" to the dataframe:

lazy val window: WindowSpec = Window.partitionBy("SERVICE_ID").orderBy("COUNTER")
val df_flag = input.filter($"FLAG" === 1)
  .withColumn("ID", row_number().over(window))
val df_other = input.filter($"FLAG" =!= 1)
  .withColumn("ID", lit(0))

// Create a group for each flag event
val df = df_flag.union(df_other)
  .withColumn("ID", max("ID").over(window.rowsBetween(Long.MinValue, 0)))
  .cache()

df.show() gives:

+----------+-------+--------+----+---+
|SERVICE_ID|COUNTER|EVENT_ID|FLAG| ID|
+----------+-------+--------+----+---+
|      1235|      1|    1111|   1|  1|
|      1235|      2|    1112|   0|  1|
|      1235|      3|    1114|   0|  1|
|      1235|      4|    2221|   1|  2|
|      1235|      5|    2225|   0|  2|
|      1235|      6|    2226|   0|  2|
|      1235|      7|    2227|   1|  3|
+----------+-------+--------+----+---+

Now that we have a column separating the events, we need to add the correct "EVENT_ID" (renamed "first_value") to each event. In addition to the "first_value", calculate and add a second column "last_value", which is the id of the next flagged event.

val df_event = df.filter($"FLAG" === 1)
  .select("EVENT_ID", "ID", "SERVICE_ID", "COUNTER")
  .withColumnRenamed("EVENT_ID", "first_value")
  .withColumn("last_value", lead($"first_value",1,0).over(window))
  .drop("COUNTER")

val df_final = df.join(df_event, Seq("ID", "SERVICE_ID"))
  .drop("ID")
  .withColumn("first_value", when($"FLAG" === 1, lit(0)).otherwise($"first_value"))

df_final.show() gives us:

+----------+-------+--------+----+-----------+----------+
|SERVICE_ID|COUNTER|EVENT_ID|FLAG|first_value|last_value|
+----------+-------+--------+----+-----------+----------+
|      1235|      1|    1111|   1|          0|      2221|
|      1235|      2|    1112|   0|       1111|      2221|
|      1235|      3|    1114|   0|       1111|      2221|
|      1235|      4|    2221|   1|          0|      2227|
|      1235|      5|    2225|   0|       2221|      2227|
|      1235|      6|    2226|   0|       2221|      2227|
|      1235|      7|    2227|   1|          0|         0|
+----------+-------+--------+----+-----------+----------+

Answer 2

Can be solved in two steps:

1. get events with "FLAG" == 1 and valid range for this event;
2. join 1. with input, by range.

Some column renaming included for visibility, can be shortened:

val window = Window.partitionBy("SERVICE_ID").orderBy("COUNTER").rowsBetween(Window.currentRow, 1)
val eventRangeDF = input.where($"FLAG" === 1)
  .withColumn("RANGE_END", max($"COUNTER").over(window))
  .withColumnRenamed("COUNTER", "RANGE_START")
  .select("SERVICE_ID", "EVENT_ID", "RANGE_START", "RANGE_END")
eventRangeDF.show(false)

val result = input.where($"FLAG" === 0).as("i").join(eventRangeDF.as("e"),
  expr("e.SERVICE_ID=i.SERVICE_ID And i.COUNTER>e.RANGE_START and i.COUNTER<e.RANGE_END"))
  .select($"i.SERVICE_ID", $"i.COUNTER", $"i.EVENT_ID", $"i.FLAG", $"e.EVENT_ID".alias("first_value"))
  // include FLAG=1
  .union(input.where($"FLAG" === 1).select($"SERVICE_ID", $"COUNTER", $"EVENT_ID", $"FLAG", lit(0).alias("first_value")))

result.sort("COUNTER").show(false)

Output:

+----------+--------+-----------+---------+
|SERVICE_ID|EVENT_ID|RANGE_START|RANGE_END|
+----------+--------+-----------+---------+
|1235      |1111    |1          |4        |
|1235      |2221    |4          |7        |
|1235      |2227    |7          |7        |
+----------+--------+-----------+---------+

+----------+-------+--------+----+-----------+
|SERVICE_ID|COUNTER|EVENT_ID|FLAG|first_value|
+----------+-------+--------+----+-----------+
|1235      |1      |1111    |1   |0          |
|1235      |2      |1112    |0   |1111       |
|1235      |3      |1114    |0   |1111       |
|1235      |4      |2221    |1   |0          |
|1235      |5      |2225    |0   |2221       |
|1235      |6      |2226    |0   |2221       |
|1235      |7      |2227    |1   |0          |
+----------+-------+--------+----+-----------+