在Spark数据框中映射

Question

Using Spark 2.x I'm making use of the dataframes.

val proposals = spark.read
  .option("header", true)
  .option("inferSchema", true)
  .option("delimiter", ";")
  .csv("/proposals.txt.gz")

proposals.printSchema()

which works fine and gives:

root
 |-- MARKETCODE: string (nullable = true)
 |-- REFDATE: string (nullable = true)
 |-- UPDTIME: string (nullable = true)
 |-- UPDTIMEMSEC: integer (nullable = true)
 |-- ENDTIME: string (nullable = true)
 |-- ENDTIMEMSEC: integer (nullable = true)
 |-- BONDCODE: string (nullable = true)

Now I'd like to calculate a time in milliseconds and thus have written a function:

def time2usecs( time:String, msec:Int )={
    val Array(hour,minute,seconds) = time.split(":").map( _.toInt )
    msec + seconds.toInt*1000 + minute.toInt*60*1000 + hour.toInt*60*60*1000
}
time2usecs( "08:13:44", 111 )


time2usecs: (time: String, msec: Int)Int
res90: Int = 29624111

The last peace of the puzzle that would be something like:

proposals.withColumn( "utime",
  proposals.select("UPDTIME","UPDTIMEMSEC")
    .map( (t,tms) => time2usecs(t,tms) ))

But I can't figure out how to do the df.select(column1, column2).map(...) part.

Answer 1

The common approach to using a method on dataframe columns in Spark is to define an UDF (User-Defined Function, see here for more information). For your case:

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

val time2usecs = udf((time: String, msec: Int) => {
  val Array(hour,minute,seconds) = time.split(":").map( _.toInt )
  msec + seconds.toInt*1000 + minute.toInt*60*1000 + hour.toInt*60*60*1000
})

val df2 = df.withColumn("utime", time2usecs($"UPDTIME", $"UPDTIMEMSEC"))

spark.implicits._ is imported here to allow the use of the $ shorthand for the col() function.

Answer 2

Why not use SQL all the way?

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

def time2usecs(time: Column, msec: Column) = {
  val bits  = split(time, ":")
  msec + bits(2).cast("int") * 1000 + bits(1).cast("int") * 60 * 1000 + 
  bits(0).cast("int") *60*60*1000
}

df.withColumn("ts", time2usecs(col(""UPDTIME"), col("UPDTIMEMSEC"))

With your code you'd have to:

proposals
  .select("UPDTIME","UPDTIMEMSEC")
  .as[(String, Int)]
  .map { case (t, s) => time2usecs(t, s) }