Spark Dataframe：访问地图功能中的下一条记录

Question

我有一个带时间戳列的DF，按此列排序。有没有办法做到这一点：对于每条记录，访问下一条记录来计算两条线之间的时间差异？我不认为这在地图功能中是可行的，因为这两行可以在不同的节点上处理。

谢谢！

Answer 1

对于Spark 1.4或更高版本，如果您可以使用Hive Context，则以下代码可能适合您：

import org.apache.spark.sql.hive.HiveContext
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.LongType
import org.apache.spark.sql._

val hc = new HiveContext(sc)
val df = hc.read.format("...").load("...")

val timestamp_column = df("timestamp_column")
val next_row_timestamp = lead(timestamp_column, 1).over(Window.orderBy(timestamp_column))

val newDF = df.withColumn("time_difference", next_row_timestamp.cast(LongType) - timestamp_column.cast(LongType))

<强>解释

在此代码中，我使用lead(e: Column, offset: Int)包（doc）中提供的functions窗口函数。此函数实际上创建了一个新列，其中e列中的数据（示例中为timestamp_column）由offset（在示例中为1）中提取。要正常工作，必须跟随over(window: WindowSpec)调用，该调用使用Window对象定义窗口。该窗口可以由分区和订单组成。在这种情况下，我只使用Window.orderBy设置订单。

最后，我使用withColumn添加一个列，其中两列之间的差异（或毫秒？不确定）与原始DataFrame的差异。

有关详细信息，以下链接可以很好地解释这个想法，并带有示例： https://databricks.com/blog/2015/07/15/introducing-window-functions-in-spark-sql.html

修改

正如评论中所指出的，上述解决方案效率非常低。作为替代方案，可以使用RDD解决方案：

val newRDD = df.rdd.zipWithIndex.flatMap {
  case (row, idx) => (0 to 1).map { lag => (idx - lag, row) }
}
.groupByKey
.values
.map { pair =>
  val pairArray = pair.toArray
  val timeDiff = {
    if (pairArray.length == 1) null
    else pairArray(1).getAs[java.sql.Timestamp]("timestamp_column").getTime - pairArray(0).getAs[java.sql.Timestamp]("timestamp_column").getTime
  }
  Row.merge(Row(timeDiff), pairArray(0))
}

val newSchema = StructType(StructField("time_diff", LongType, true) +: df.schema.fields)
val newDf = df.sqlContext.createDataFrame(newRDD, newSchema)

newDF中的结果数据框将有一个新列“time_diff”，其中包含当前行与下一行之间的时间差（以毫秒为单位）。