在 Dataframe 列中添加函数的结果 [Spark Scala]

Question

我想做一些计算并将其添加到现有数据框中。我有以下函数根据经度和纬度计算地址space。

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

我创建了一个具有以下架构的数据框：

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

我想 add/append 这个 Dataframe 的新列 H3Address，其中地址 space 是根据 lat 和 [=该行的 18=]。

这是我想要实现的数据框的一小部分：

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

我试过类似的方法：

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

但这没有用。

有人可以帮我吗？

编辑：

添加@Garib的建议后，得到如下几行：

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

如何解决这个错误？尝试了很多东西。例如通过使用 lit() 函数

编辑2:

使用正确的lit()方法后，提出的方案奏效了。

解决方法： df.withColumn("H3Address", getH3Address(col("lat"), col("lon"), lit(10)))

Answer 1

您应该从您的函数中创建一个 UDF。

User-Defined Functions (UDFs) are user-programmable routines that act on one row

例如：

val getH3Address = udf(
  // write here the logic of your function. I used a dummy logic (x+y) just for this example.
  (x: Double, y: Double) => {
    (x + y).toString
  })

val df = Seq((1, 2, "aa"), (2, 3, "bb"), (3, 4, "cc")).toDF("lat", "lon", "value")
df.withColumn("H3Address", getH3Address(col("lat"), col("lon"))).show()

您可以在此处阅读有关 UDF 的更多信息： https://spark.apache.org/docs/latest/sql-ref-functions-udf-scalar.html

在 Dataframe 列中添加函数的结果 [Spark Scala]

Adding the result of a function in a Dataframe column [Spark Scala]

scala

dataframe

apache-spark

user-defined-functions

calculation