Scala 如何在对行执行映射操作后保留Spark数据帧的列结构

我试图对Spark数据帧的每一行应用一个函数，如示例所示

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()

---

该函数起作用了，但我现在有一列的条目是列表，而不是3列的int。我想要回我的命名列。

您可以定义一个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|
+---+---+---+

---

<>而不是将行文件转换成行，考虑使用内置的火花API函数<代码> //否则< /代码>，如下：

import org.apache.spark.sql.functions._
import spark.implicits._

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

val vFrom = 0
val vTo = 5

val cols = df.columns  // Filter for specific columns if necessary

df.select( cols.map( c => 
    when(col(c) === vFrom, vTo).otherwise(col(c)).as(c)
  ): _*
).show
// +---+---+---+
// |  x|  y|  z|
// +---+---+---+
// |  1|  2|  5|
// |  5|  5|  1|
// |  5|  5|  5|
// +---+---+---+

---

有多种方法可以做到这一点，以下是一些：

df.map(row => {
      val size = row.size
      var seq: Seq[Int] = Seq.empty[Int]
      for (a <- 0 to size - 1) {
        val value: Int = row(a).asInstanceOf[Int]
        val newVal: Int = value match {
          case 0 =>
            5
          case _ =>
            value
        }
        seq = seq :+ newVal
      }
      Row.fromSeq(seq)
    })(RowEncoder.apply(df.schema))

---

这些都是很好的答案，但我接受了这个答案，因为第一个解决方案比我的玩具示例更适用于行以更复杂的方式转换的情况。不确定为什么有人否决了一个得到3个有趣答案的问题。这永远不会发生在交叉验证。投票-事实上，它经常发生，这总是一个很好的答案。

df.map(row => {
      val size = row.size
      var seq: Seq[Int] = Seq.empty[Int]
      for (a <- 0 to size - 1) {
        val value: Int = row(a).asInstanceOf[Int]
        val newVal: Int = value match {
          case 0 =>
            5
          case _ =>
            value
        }
        seq = seq :+ newVal
      }
      Row.fromSeq(seq)
    })(RowEncoder.apply(df.schema))

 val columns = df.columns
    df.select(
        columns.map(c => when(col(c) === 0, 5).otherwise(col(c)).as(c)): _*)
      .show()

def fun: (Int => Int) = { x =>
      if (x == 0) 5 else x
    }
    val function = udf(fun)
    df.select(function(col("x")).as("x"),
              function(col("y")).as("y"),
              function(col("z")).as("z"))
      .show()

def checkZero(a: Int): Int = if (a == 0) 5 else a

      df.map {
        case Row(a: Int, b: Int, c: Int) =>
          Row(checkZero(a), checkZero(b), checkZero(c))
      } { RowEncoder.apply(df.schema) }
      .show()