我应该为 Spark Dataframe udf 中的元组使用什么数据类型？

Question

Input:

val df = Seq((10, (35, 25))).toDF("id", "scorePair")
df.show
+---+---------+
| id|scorePair|
+---+---------+
| 10| [35, 25]|
+---+---------+

Expected output:

+---+-----------+
| id|totalScore |
+---+-----------+
| 10|         60|
+---+-----------+

Wanted to do something like this, but it does not accept Row type:

// error
val add = udf((row: Row) => {row match {case (a: Int, b: Int) => a + b}})
df.withColumn("totalScore", add(col("scorePair")))

Why Row type is not correct thinking of

"Dataframe is an alias for Dataset[Row]"

?

What type should I use? How can I achieve it?

---

• I emphasize on the type Row , because at lest I manage to use Row in the following way (which treats each cell of a column as a Row ) to achieve:

val add = udf((rows: Seq[Row]) => {rows.map {case Row(a: Int, b: Int) => a + b}})
df.groupBy("id").agg(collect_list("scorePair") as "pairSeq").withColumn("totalScore1", add(col("pairSeq"))).select(col("id"), explode(col("totalScore1")) as "totalScore").show
+---+----------+
| id|totalScore|
+---+----------+
| 10|        60|
+---+----------+

But that's really not clean!

Answer 1

You can use either row.getAs[Int](0) , row.get(0).asInstanceOf[Int] , row.getInt(0) to get the value from row

val df = Seq(
  (10, (35, 25))
).toDF("id", "scorePair")


val add = udf((row: Row) => {row.getInt(0) + row.getInt(1)})

df.withColumn("totalScore", add($"scorePair")).show(false)

df.select($"id", $"scorePair._1" + $"scorePair._2" as "totalScore").show(false)

Output:

+---+----------+
|id |totalScore|
+---+----------+
|10 |60        |
+---+----------+

Answer 2

The aggregate function is the easiest way to sum all the numbers in an ArrayType column. This post has a full example. Here's the snippet:

val resDF = df.withColumn(
  "totalScore",
  aggregate(
    col("scorePair"),
    lit(0),
    (col1: Column, col2: Column) => col1 + col2
  )
)

You want to avoid UDFs whenever possible. This solution only works for Spark 3+.