Spark ML中无法explode拆分PCA输出Vector列绘制散点图

问题说明 #

使用Spark ML库PCA算法计算得到的结果中，pca列为org.apache.spark.ml.linalg.Vector类型，存储2个主成分维度值，原始表结构如下：

_________
id | pca
__________
1  |[0.1,0.1]
2  |[0.2,0.2]
3  |[0.4,0.4]
4  |[0.3,0.3]
...

需要拆分出独立的X、Y维度列，适配散点图绘制需求，目标表结构如下：

_________
id | X | Y
__________
1  |0.1|0.1
2  |0.2|0.2
3  |0.4|0.4
4  |0.3|0.3
...

直接使用col("pca")[0]、col("pca")[1]下标方式取值会失败，原因是Spark SQL原生下标语法仅支持ArrayType类型列，MLlib的Vector是自定义结构体类型，不属于原生数组，无法直接通过下标访问。
原有实现代码如下：

import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.feature.VectorAssembler


val convertToVector = udf((array: Array[Double]) => {
  Vectors.dense(array.toArray)
})

val convertToDouble = udf((array: Array[Float]) => {
  array.map(_.toDouble).toArray
})

val ds = model.userFactors.withColumn("features", convertToDouble($"features"))
val userMatrixDs = ds.withColumn("features", convertToVector($"features"))

val pca = new PCA()
    .setInputCol("features")
    .setOutputCol("pca")
    .setK(2)
    .fit(userMatrixDs)
val result = pca.transform(userMatrixDs).select("id","pca");

result.show()

// 以下直接索引Vector的写法会运行失败
result.select(
    result.id,
    result.col("pca")[0].as("eigenVector1"),
    result.col("pca")[1].as("eigenVector2")
  )
  .show()

解决方案 #

方法1：内置函数转换（推荐，Spark 2.4+版本支持，性能最优） #

使用Spark ML内置的vector_to_array函数先将Vector列转换为原生Double数组，再通过下标取值即可，无需自定义UDF，经过Spark原生优化性能更好。
代码实现：

// 导入内置转换函数
import org.apache.spark.ml.linalg.vector_to_array

val splitResult = result
  // 先将Vector列转为数组列
  .withColumn("pca_arr", vector_to_array($"pca"))
  .select(
    $"id",
    $"pca_arr"(0).as("X"),
    $"pca_arr"(1).as("Y")
  )

// 打印验证结果
splitResult.show()

方法2：自定义UDF取值（兼容所有Spark版本） #

如果使用低于2.4的Spark版本，没有内置转换函数，可以自定义UDF提取Vector指定位置的元素：

import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.functions.lit

// 定义UDF：输入Vector和索引位置，返回对应位置的Double值
val getVectorElem = udf((vec: org.apache.spark.ml.linalg.Vector, index: Int) => vec(index))

val splitResult = result.select(
  $"id",
  getVectorElem($"pca", lit(0)).as("X"),
  getVectorElem($"pca", lit(1)).as("Y")
)

// 打印验证结果
splitResult.show()

注意事项 #

• 不要混用org.apache.spark.ml.linalg.Vector和org.apache.spark.mllib.linalg.Vector两种类型，自定义UDF的参数类型需要和PCA输出的Vector类型保持一致，否则会报类型不匹配错误。
• 拆分得到的X、Y列为原生Double类型，可直接导出为CSV/JSON等格式，对接可视化工具绘制散点图。

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