我有一个由两个双打阵列组成的数据帧.我想创建一个新列,它是将欧几里德距离函数应用于前两列的结果,即如果我有：
A B (1,2) (1,3) (2,3) (3,4)
创建：
A B C (1,3) 1 (2,4) 1.4
我的数据架构是：
df.schema.foreach(println) StructField(col1,ArrayType(DoubleType,false),false) StructField(col2,true)
def distance(xs: Array[Double],ys: Array[Double]) = {
sqrt((xs zip ys).map { case (x,y) => pow(y - x,2) }.sum)
}
我收到类型错误：
df.withColumn("distances",distance($"col1",$"col2"))
<console>:68: error: type mismatch;
found : org.apache.spark.sql.ColumnName
required: Array[Double]
ids_with_predictions_centroids3.withColumn("distances",$"col2"))
我知道我必须遍历每列的元素,但我无法找到如何在任何地方执行此操作的说明.我是Scala编程的新手.
解决方法
要在数据帧上使用自定义函数,需要将其定义为UDF.例如,这可以完成如下：
val distance = udf((xs: WrappedArray[Double],ys: WrappedArray[Double]) => {
math.sqrt((xs zip ys).map { case (x,y) => math.pow(y - x,2) }.sum)
})
df.withColumn("C",distance($"A",$"B")).show()
请注意,此处需要使用WrappedArray(或Seq).
结果数据帧：
+----------+----------+------------------+ | A| B| C| +----------+----------+------------------+ |[1.0,2.0]|[1.0,3.0]| 1.0| |[2.0,3.0]|[3.0,4.0]|1.4142135623730951| +----------+----------+------------------+
