如果您对使用私有API还满意,那么我发现了一种有效的方法:将整个自引用结构视为用户定义的类型。我遵循这个答案:
https://stackoverflow.com/a/51957666/1823254
.
package org.apache.spark.custom.udts // we're calling some private API so need to be under 'org.apache.spark'
import java.io._
import org.apache.spark.sql.types.{DataType, UDTRegistration, UserDefinedType}
class BranchUDT extends UserDefinedType[Branch] {
override def sqlType: DataType = org.apache.spark.sql.types.BinaryType
override def serialize(obj: Branch): Any = {
val bos = new ByteArrayOutputStream()
val oos = new ObjectOutputStream(bos)
oos.writeObject(obj)
bos.toByteArray
}
override def deserialize(datum: Any): Branch = {
val bis = new ByteArrayInputStream(datum.asInstanceOf[Array[Byte]])
val ois = new ObjectInputStream(bis)
val obj = ois.readObject()
obj.asInstanceOf[Branch]
}
override def userClass: Class[Branch] = classOf[Branch]
}
object BranchUDT {
def register() = UDTRegistration.register(classOf[Branch].getName, classOf[BranchUDT].getName)
}
BranchUDT.register()
val trees = Seq(Tree(1, List(Branch(2, List.empty), Branch(3, List(Branch(4, List.empty))))))
val ds = spark.createDataset(trees)
ds.show(false)
//+---+----------------------------------------------------+
//|id |branches |
//+---+----------------------------------------------------+
//|1 |[Branch(2,List()), Branch(3,List(Branch(4,List())))]|
//+---+----------------------------------------------------+