Spark Scala如何将H3地址计算函数结果写入DataFrame新列

Spark DataFrame 新增H3地址列实现方案 #

需求说明 #

需要基于现有DataFrame的经纬度字段，计算得到H3地址并新增为H3Address列，已定义基础计算函数如下：

def getH3Address(x: Double, y: Double): String ={
    h3.get.geoToH3Address(x,y)
}

现有DataFrame的Schema结构：

root
 |-- lat: double (nullable = true)
 |-- lon: double (nullable = true)
 |-- elevation: integer (nullable = true)

预期输出效果示例：

+----+------------------+---------+---------+
| lat|               lon|elevation|H3Address|
+----+------------------+---------+---------+
|51.0|               3.0|       13|   a3af83|
|51.0| 3.000277777777778|       13|   a3zf83|
|51.0|3.0005555555555556|       12|   a1qf82|
|51.0|3.0008333333333335|       12|   l3xf83|

常见错误及解决 #

错误1：直接调用普通函数传列参数

最初写法如下，运行失败：

df.withColumn("H3Address", geoToH3Address(df.select(df("lat")), df.select(df("lon")))

错误原因：
普通Scala函数无法直接处理Spark的列类型，withColumn方法的第二个参数要求返回org.apache.spark.sql.Column类型，需要先将自定义函数封装为Spark UDF才可以操作DataFrame的行级数据。

错误2：UDF传参直接使用字面量

封装UDF后代码如下，运行报类型不匹配错误：

val getH3Address = udf(
      (lat: Double, lon: Double, res: Int) => {
        h3.get.geoToH3Address(lat,lon,res).toString
      })
var res : Int = 10

val DF_edit = df.withColumn("H3Address", 
 getH3Address(col("lat"), col("lon"), 10))

报错信息：

[error]  type mismatch;
  found   : Int
  required: org.apache.spark.sql.Column

错误原因：
UDF的所有入参都必须是Spark Column类型，常量字面量需要用lit()函数包装为Column对象才能传入。

最终正确实现 #

修正后代码如下：

// 导入lit函数依赖
import org.apache.spark.sql.functions.lit

val getH3Address = udf(
      (lat: Double, lon: Double, res: Int) => {
        h3.get.geoToH3Address(lat,lon,res).toString
      })

val DF_edit = df.withColumn("H3Address", getH3Address(col("lat"), col("lon"), lit(10)))

内容的提问来源于stack exchange，提问作者EnesK