如何将Spark DataFrame中EmbeddedList类型列转为Array<String>

一、转换EmbeddedList为Spark原生Array #

EmbeddedList是OrientDB自定义的UDT，内部封装了数组结构，可通过两种方式转换为Spark原生Array<String>：

方法1：直接提取UDT内部字段 #

若EmbeddedList类包含values属性（对应UDT定义的存储字段），可直接在DataFrame中引用该字段提取数组：

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

// 提取EmbeddedList内部数组并转为Array<String>
val dfWithArray = df.withColumn("title_array", col("title.values").cast("array<string>"))

方法2：自定义UDF转换 #

如果直接提取字段无效，可编写自定义UDF解析EmbeddedList：

import org.apache.spark.sql.functions.udf
import com.orientechnologies.spark.udts.EmbeddedList

// 定义UDF：将EmbeddedList转为Array<String>
val embeddedListToArray = udf((el: EmbeddedList) => el.values.asInstanceOf[Array[String]])

// 应用UDF生成新列
val dfWithArray = df.withColumn("title_array", embeddedListToArray(col("title")))

二、修正数组拼接逻辑（解决concat报错） #

你遇到的concat报错是因为Spark的concat函数仅支持字符串/二进制类型拼接，不支持数组合并。要合并character（单个字符串）和title_array（字符串数组），可使用flatten函数将嵌套数组扁平化：

val vertices = dfWithArray
  .select(explode(flatten(array(array(col("character")), col("title_array")))) as "x")
  .distinct()
  .rdd.map(_.getAs[String]("x"))
  .zipWithIndex()
  .map(_.swap)

替代方案：用Union实现 #

如果不需要合并数组，也可分别提取字符和标题元素后再合并：

val chars = df.select(col("character").as("x"))
val titles = df.select(explode(col("title_array")).as("x"))

val vertices = chars.union(titles)
  .distinct()
  .rdd.map(_.getAs[String]("x"))
  .zipWithIndex()
  .map(_.swap)

完整示例代码 #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import com.orientechnologies.spark.udts.EmbeddedList

object OrientDBConversion {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("OrientDBTest").master("local[*]").getOrCreate()
    import spark.implicits._

    // 模拟输入数据
    val df = Seq(
      "Tony Stark" -> EmbeddedList(Array("Iron Man")),
      "James Buchanan Barnes" -> EmbeddedList(Array("Captain America: The First Avenger", "Captain America: The Winter Soldier", "Captain America: Civil War", "Avengers: Infinity War")),
      "Marcus Bledsoe" -> EmbeddedList(Array("Captain America: The Winter Soldier"))
    ).toDF("character", "title")

    // 转换EmbeddedList为Array<String>
    val embeddedListToArray = udf((el: EmbeddedList) => el.values.asInstanceOf[Array[String]])
    val dfWithArray = df.withColumn("title_array", embeddedListToArray(col("title")))

    // 生成目标RDD
    val vertices = dfWithArray
      .select(explode(flatten(array(array(col("character")), col("title_array")))) as "x")
      .distinct()
      .rdd.map(_.getAs[String]("x"))
      .zipWithIndex()
      .map(_.swap)

    // 输出结果
    vertices.foreach(println)
    spark.stop()
  }
}

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