在Spark Scala中将ArrayType列传递给UDF

Question

我在Scala的Spark数据框架中有一列，它是使用

 agg(collect_list(struct(col(abc), col(aaa)).as(def)

我想将此列传递给UDF，以便进一步处理，以处理此聚合列中的一个索引。

当我将参数传递给我的UDF时：

.withColumn(def, remove
            (col(xyz), col(def)))

UDF-类型为Seq[Row]：val removeUnstableActivations:UserDefinedFunction=UDF（（xyz:java.util.Date，def:Seq[Row]）

我收到错误：

Exception encountered when invoking run on a nested suite - Schema for type org.apache.spark.sql.Row is not supported

我应该如何传递这些列，UDF中列的数据类型应该是什么？

Answer 1

事实上，不支持类型Row的模式，但您可以返回案例类。Spark将返回的案例类视为结构类型。例如：

import org.apache.spark.sql.functions._
import spark.implicits._
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.Row

val df = Seq(
  (1, "a"),
  (2, "b"),
  (3, "c")
).toDF("number", "word")

val aggDf = df.agg(
  collect_list(struct(col("number"), col("word"))) as "aggColumn"
)

aggDf.printSchema()
// |-- aggColumn: array (nullable = true)
// |    |-- element: struct (containsNull = true)
// |    |    |-- number: string (nullable = true)
// |    |    |-- word: integer (nullable = false)

case class ReturnSchema(word: String, number: Int)

val myUdf: UserDefinedFunction =
  udf((collection: Seq[Row]) => {
    collection.map(r => {
      val word   = r.getAs[String]("word")
      val newNumber = r.getAs[Int]("number") * 100

      new ReturnSchema(word, newNumber)
    })
  })
  
val finalDf = aggDf.select(myUdf(col("aggColumn")).as("udfTranformedColumn"))

finalDf.printSchema
// root
//  |-- udfTranformedColumn: array (nullable = true)
//  |    |-- element: struct (containsNull = true)
//  |    |    |-- word: string (nullable = true)
//  |    |    |-- number: integer (nullable = false)

finalDf.show(false)
// +------------------------------+
// |udfTranformedColumn           |
// +------------------------------+
// |[[a, 100], [b, 200], [c, 300]]|
// +------------------------------+