scala - 使用 flatmap 将一行 spark 数据集分解为多行并添加列

Question

我有一个具有以下架构的 DataFrame :

root
 |-- journal: string (nullable = true)
 |-- topicDistribution: vector (nullable = true)

topicDistribution 字段是 double 向量:[0.1, 0.2 0.15 ...]

我想要的是将每一行分解成几行以获得以下架构:

root
 |-- journal: string
 |-- topic-prob: double // this is the value from the vector
 |-- topic-id : integer // this is the index of the value from the vector

为了澄清，我创建了一个案例类:

case class JournalDis(journal: String, topic_id: Integer, prob: Double)

我已经设法以一种非常笨拙的方式使用 dataset.explode 实现了这一点:

val df1 = df.explode("topicDistribution", "topic") {
    topics: DenseVector => topics.toArray.zipWithIndex
}.select("journal", "topic")
val df2 = df1.withColumn("topic_id", df1("topic").getItem("_2")).withColumn("topic_prob", df1("topic").getItem("_1")).drop(df1("topic"))

但dataset.explode 已弃用。我想知道如何使用 flatmap 方法实现这一点？

Answer 1

未测试但应该可以工作:

import spark.implicits._
import org.apache.spark.ml.linalg.Vector

df.as[(String, Vector)].flatMap { 
  case (j, ps) => ps.toArray.zipWithIndex.map { 
    case (p, ti) => JournalDis(j, ti, p)
  }
}