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

时间:2019-05-02 14:42:19

标签: scala apache-spark apache-spark-sql

我正在尝试将功能应用于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列的Ints。我想要我的命名列。

3 个答案:

答案 0 :(得分:3)

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

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

然后,您可以将UDF应用于特定的列,或者,如果需要,可以将DFD应用于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|
+---+---+---+

答案 1 :(得分:2)

考虑按以下方式使用内置的Spark API函数when/otherwise,而不是按行转换DataFrame:

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|
// +---+---+---+

答案 2 :(得分:1)

这里有多种实现方法:

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