scala - 如何在行上执行映射操作后保留 Spark Dataframe 的列结构

Question

我正在尝试将函数应用于 Spark DataFrame 的每一行，如示例中所示。

val df = sc.parallelize(
  Seq((1, 2, 0), (0, 0, 1), (0, 0, 0))).toDF("x", "y", "z")
df.show()

这产生

+---+---+---+
|  x|  y|  z|
+---+---+---+
|  1|  2|  0|
|  0|  0|  1|
|  0|  0|  0|
+---+---+---+

假设我想对每一行中的值做一些事情，例如将 0 更改为 5。

val b = df.map(row => row.toSeq.map(x => x match{
    case 0 => 5
    case x: Int => x
}))

b.show()

+---------+
|    value|
+---------+
|[1, 2, 5]|
|[5, 5, 1]|
|[5, 5, 5]|
+---------+

该函数有效，但我现在有一列，其条目是列表，而不是 3 列 Int。我想要我的命名列回来。

Answer 1

您可以定义一个 UDF 来应用此替换。例如:

def subsDef(k: Int): Int = if(k==0) 5  else k
val subs = udf[Int, Int](subsDef)

然后，您可以将 UDF 应用于特定列，或者，如果您愿意，也可以应用于 DF 的每一列:

// to a single column, for example "x"
df = df.withColumn("x", subs(col("x")))
df.show()
+---+---+---+
|  x|  y|  z|
+---+---+---+
|  1|  2|  0|
|  5|  0|  1|
|  5|  0|  0|
+---+---+---+



// to every columns of DF
df.columns.foreach(c => {
      df = df.withColumn(c, subs(col(c)))
    })
df.show()
+---+---+---+
|  x|  y|  z|
+---+---+---+
|  1|  2|  5|
|  5|  5|  1|
|  5|  5|  5|
+---+---+---+