Spark Dataframe:访问映射函数中的下一条记录

Question

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

谢谢！

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))

解释：

在这段代码中，我使用的是functions包（doc）中提供的lead(e: Column, offset: Int)窗口函数。此函数实际上创建了一个新列，其中列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"，其中包含当前行和下一行之间的时间差（以毫秒为单位）。